diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index d116f8ae2925..6c2f5153a4ea 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -13,12 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import ctypes +import errno import functools import os import platform import sys +import time from dataclasses import dataclass -from typing import List, Optional, Union +from enum import Enum +from typing import Any, List, Optional, Protocol, Union import pynvml import torch @@ -33,6 +36,141 @@ from .logger import logger from .mapping import Mapping +_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S = 20.0 +_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S = 0.01 +_MNNVL_CHECKPOINT_REQUEST_CLEANUP_TIMEOUT_S = 0.1 +_MNNVL_CHECKPOINT_ALLGATHER_TAG = 31415 +_MNNVL_CHECKPOINT_ORPHANED_REQUESTS: list[Any] = [] + + +class MnnvlCheckpointCommunicator(Protocol): + """Structural contract required by bounded MNNVL checkpoint collectives.""" + + def Get_rank(self) -> int: + """Return this process's rank in the communicator.""" + ... + + def Get_size(self) -> int: + """Return the communicator size.""" + ... + + def barrier(self) -> None: + """Synchronize all communicator members.""" + ... + + def allgather(self, value: Any) -> list[Any]: + """Gather a Python object from every communicator member.""" + ... + + def isend(self, value: Any, *, dest: int, tag: int) -> Any: + """Start a nonblocking Python-object send.""" + ... + + def irecv(self, *, source: int, tag: int) -> Any: + """Start a nonblocking Python-object receive.""" + ... + + +def _cancel_checkpoint_requests(requests: list[Any]) -> None: + """Bound cancellation cleanup so timeout handling cannot hang in MPI.""" + for request in requests: + try: + request.Cancel() + except Exception as error: + logger.warning(f"Failed to cancel MNNVL checkpoint request: {error}") + + pending_requests = list(requests) + cleanup_timeout_s = min( + _MNNVL_CHECKPOINT_REQUEST_CLEANUP_TIMEOUT_S, + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S, + ) + deadline = time.monotonic() + cleanup_timeout_s + while pending_requests: + incomplete_requests = [] + for request in pending_requests: + try: + ready, _ = request.test() + except Exception as error: + logger.warning(f"Failed to poll MNNVL checkpoint request cleanup: {error}") + incomplete_requests.append(request) + continue + if not ready: + incomplete_requests.append(request) + pending_requests = incomplete_requests + if not pending_requests or time.monotonic() >= deadline: + break + time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + + if pending_requests: + # An active receive must not be freed: mpi4py owns its receive buffer, + # and MPI may still write into it. Keep the wrappers alive until the + # enclosing fail-closed path terminates the worker. + _MNNVL_CHECKPOINT_ORPHANED_REQUESTS.extend(pending_requests) + logger.error( + f"Retaining {len(pending_requests)} active MNNVL checkpoint requests " + "until worker termination" + ) + + +def _checkpoint_allgather( + comm: MnnvlCheckpointCommunicator, + value: Any, + *, + operation: str, +) -> list[Any]: + """Run a bounded object allgather over nonblocking point-to-point requests. + + Production mpi4py communicators provide ``isend`` and ``irecv`` but no + nonblocking object allgather. The blocking fallback is retained only for + lightweight test communicators that do not implement point-to-point APIs. + """ + isend = getattr(comm, "isend", None) + irecv = getattr(comm, "irecv", None) + if isend is None or irecv is None: + return comm.allgather(value) + + rank = comm.Get_rank() + size = comm.Get_size() + results: list[Any] = [None] * size + results[rank] = value + receive_requests: dict[int, Any] = {} + send_requests: list[Any] = [] + try: + for peer in range(size): + if peer != rank: + receive_requests[peer] = irecv( + source=peer, + tag=_MNNVL_CHECKPOINT_ALLGATHER_TAG, + ) + for peer in range(size): + if peer != rank: + send_requests.append( + isend( + value, + dest=peer, + tag=_MNNVL_CHECKPOINT_ALLGATHER_TAG, + ) + ) + if not receive_requests: + return results + deadline = time.monotonic() + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S + while receive_requests or send_requests: + for peer, request in list(receive_requests.items()): + ready, result = request.test() + if ready: + results[peer] = result + del receive_requests[peer] + send_requests = [request for request in send_requests if not request.test()[0]] + if not receive_requests and not send_requests: + return results + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for MNNVL checkpoint {operation} allgather") + time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + return results + except Exception: + _cancel_checkpoint_requests([*receive_requests.values(), *send_requests]) + raise + def _check_cu_result(cu_func_ret): if isinstance(cu_func_ret, tuple): @@ -51,6 +189,29 @@ 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 + state: _MnnvlAllocationState = _MnnvlAllocationState.MAPPED + pending_comm: Any = None + + class MnnvlMemory: """MNNVL memory management for tensor parallel (TP) operations.""" @@ -101,6 +262,11 @@ def __del__(self): if hasattr(self, "ptr"): type(self).close_mnnvl_memory(self.ptr) + @property + def mapped(self) -> bool: + """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() return pack_strided_memory( @@ -110,7 +276,7 @@ def as_torch_strided_tensor(self, dtype): @property def local_mem_handle(self) -> int: """Return the local rank's CUmemGenericAllocationHandle.""" - _, _, mem_handles, _, _, _ = type(self).allocated_map[self.ptr] + mem_handles = type(self).allocated_map[self.ptr].mem_handles comm_rank = type(self).comm.Get_rank() return int(mem_handles[comm_rank]) @@ -195,6 +361,180 @@ def new_mnnvl_memory_address(cls, mapping: Mapping, size: int): cls.current_rank_stride = current_rank_stride cls.current_mem_offset = 0 + @classmethod + def _create_and_map_handles( + cls, + comm, + aligned_size: int, + start_address: int, + rank_stride: int, + address_offset: int, + ) -> List[Any]: + local_handle = None + exported_handle = None + pidfds = [] + remote_fds = [] + mem_handles = [None] * comm.Get_size() + mapped_rank_ptrs = [] + is_fabric = False + try: + local_error = None + local_handle_data = None + local_pid = None + try: + dev_id = int(_check_cu_result(cuda.cuCtxGetDevice())) + assert dev_id == MnnvlMemory.dev_id, ( + f"Different dev_id found dev_id={dev_id} but " + f"MnnvlMemory.dev_id={MnnvlMemory.dev_id}" + ) + allocation_prop = MnnvlMemory.get_allocation_prop(dev_id) + is_fabric = ( + allocation_prop.requestedHandleTypes + == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ) + local_handle = _check_cu_result( + cuda.cuMemCreate(aligned_size, allocation_prop, flags=0) + ) + exported_handle = _check_cu_result( + cuda.cuMemExportToShareableHandle( + local_handle, allocation_prop.requestedHandleTypes, 0 + ) + ) + local_handle_data = exported_handle.data if is_fabric else int(exported_handle) + local_pid = os.getpid() + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + + exported_by_rank = _checkpoint_allgather( + comm, + { + "error": local_error, + "handle": local_handle_data, + "is_fabric": is_fabric, + "pid": local_pid, + }, + operation="handle export", + ) + export_errors = [ + f"rank {rank}: {payload['error']}" + for rank, payload in enumerate(exported_by_rank) + if payload["error"] is not None + ] + if export_errors: + raise RuntimeError( + "MNNVL handle export failed before mapping:\n" + "\n".join(export_errors) + ) + fabric_modes = {payload["is_fabric"] for payload in exported_by_rank} + if len(fabric_modes) != 1: + raise RuntimeError("MNNVL ranks selected inconsistent shareable handle types") + is_fabric = fabric_modes.pop() + if is_fabric: + all_handles_data = [payload["handle"] for payload in exported_by_rank] + else: + all_exported_fds = [payload["handle"] for payload in exported_by_rank] + all_pids = [payload["pid"] for payload in exported_by_rank] + syscall = ctypes.CDLL(None, use_errno=True).syscall + fd_import_error = None + try: + for pid in all_pids: + pidfd = syscall(434, pid, 0) + if pidfd < 0: + err = ctypes.get_errno() + raise RuntimeError( + f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" + ) + pidfds.append(pidfd) + for pidfd, fd in zip(pidfds, all_exported_fds): + remote_fd = syscall(438, pidfd, fd, 0) + if remote_fd < 0: + err = ctypes.get_errno() + error_msg = ( + f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno " + f"{err}: {os.strerror(err)}." + ) + if err == errno.EPERM: + error_msg += ( + " Permission denied. If running in a container, try adding " + "--cap-add=SYS_PTRACE to your docker run command." + ) + elif err == errno.ENOSYS: + error_msg += ( + " This may be due to kernel version (requires Linux 5.6+)." + ) + raise RuntimeError(error_msg) + remote_fds.append(remote_fd) + except Exception as error: + fd_import_error = f"{type(error).__name__}: {error}" + + fd_import_errors = _checkpoint_allgather( + comm, + fd_import_error, + operation="POSIX file descriptor import readiness", + ) + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(fd_import_errors) + if error is not None + ] + if failed_ranks: + raise RuntimeError( + "MNNVL POSIX file descriptor import failed on one or more ranks:\n" + + "\n".join(failed_ranks) + ) + all_handles_data = remote_fds + + access_desc = cuda.CUmemAccessDesc() + access_desc.location = allocation_prop.location + access_desc.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + for rank, handle_data in enumerate(all_handles_data): + rank_ptr = start_address + rank_stride * rank + address_offset + handle = ( + local_handle + if rank == comm.Get_rank() + else _check_cu_result( + cuda.cuMemImportFromShareableHandle( + handle_data, allocation_prop.requestedHandleTypes + ) + ) + ) + mem_handles[rank] = handle + _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, handle, 0)) + mapped_rank_ptrs.append(rank_ptr) + _check_cu_result(cuda.cuMemSetAccess(rank_ptr, aligned_size, [access_desc], 1)) + return mem_handles + except Exception: + for rank_ptr in reversed(mapped_rank_ptrs): + try: + _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) + except RuntimeError as error: + logger.warning(f"Failed to unmap incomplete MNNVL allocation: {error}") + + handles_to_release = [handle for handle in mem_handles if handle is not None] + if local_handle is not None and mem_handles[comm.Get_rank()] is None: + handles_to_release.append(local_handle) + for handle in handles_to_release: + try: + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning(f"Failed to release incomplete MNNVL allocation: {error}") + raise + finally: + for pidfd in pidfds: + try: + os.close(pidfd) + except OSError as error: + logger.warning(f"Failed to close MNNVL pidfd: {error}") + for remote_fd in remote_fds: + try: + os.close(remote_fd) + except OSError as error: + logger.warning(f"Failed to close imported MNNVL file descriptor: {error}") + if not is_fabric and exported_handle is not None: + try: + os.close(int(exported_handle)) + except OSError as error: + logger.warning(f"Failed to close exported MNNVL file descriptor: {error}") + @classmethod def open_mnnvl_memory(cls, mapping: Mapping, size: int): # Ensure MnnvlMemory is initialized (for dev_id and allocation_granularity) @@ -210,6 +550,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." @@ -227,133 +573,23 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): assert cls.current_mem_offset + aligned_size <= cls.current_rank_stride - allocation_prop = cls.get_allocation_prop(dev_id) - allocated_mem_handle = _check_cu_result( - cuda.cuMemCreate(aligned_size, allocation_prop, flags=0) - ) - exported_fabric_handle = _check_cu_result( - cuda.cuMemExportToShareableHandle( - allocated_mem_handle, allocation_prop.requestedHandleTypes, 0 - ) - ) - pidfds = [] - remote_fds = [] try: - if ( - allocation_prop.requestedHandleTypes - == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC - ): - all_handles_data = comm.allgather(exported_fabric_handle.data) - else: - all_handles_data = comm.allgather(exported_fabric_handle) - all_pids = comm.allgather(os.getpid()) - libc = ctypes.CDLL(None, use_errno=True) - syscall = libc.syscall - SYS_pidfd_open = 434 - SYS_pidfd_getfd = 438 - for i, pid in enumerate(all_pids): - pidfd = syscall(SYS_pidfd_open, pid, 0) - if pidfd < 0: - err = ctypes.get_errno() - raise RuntimeError( - f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" - ) - pidfds.append(pidfd) - - for i, (pidfd, fd) in enumerate(zip(pidfds, all_handles_data)): - remote_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0) - if remote_fd < 0: - err = ctypes.get_errno() - error_msg = f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno {err}: {os.strerror(err)}." - if err == 1: # EPERM - error_msg += ( - " Permission denied. If running in a container, try adding --cap-add=SYS_PTRACE " - "to your docker run command." - ) - else: - error_msg += " This may be due to kernel version (requires Linux 5.6+)." - raise RuntimeError(error_msg) - remote_fds.append(remote_fd) - - all_handles_data = remote_fds - except Exception: - # Release resources on failure path to avoid leaks; then re-raise. - if isinstance(exported_fabric_handle, int): - try: - os.close(exported_fabric_handle) - except OSError as e: - logger.warning( - "Failed to close exported shareable handle on error: %s", - e, - ) - try: - _check_cu_result(cuda.cuMemRelease(allocated_mem_handle)) - except RuntimeError as e: - logger.warning( - "cuMemRelease failed during error cleanup (original error will be raised): %s", - e, - ) - for _pidfd in pidfds: - try: - os.close(_pidfd) - except OSError: - pass - for _rfd in remote_fds: - try: - os.close(_rfd) - except OSError: - pass - raise - # all_handles_data like b'\x00\x00\x00 \x00\x00\x00\x00\x8f\xec\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\t\x00\x00\x00\x00\x00\x1d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' # noqa: E501 - # can use buf = memoryview(data) to import if using plain buffer for data. - - madesc = cuda.CUmemAccessDesc() - madesc.location = allocation_prop.location - madesc.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE - - mem_handles = [None] * comm_size - mem_handles[comm_rank] = allocated_mem_handle - mapped_rank_ptrs = [] - - try: - for i, remote_handle_data in enumerate(all_handles_data): - rank_ptr = ( - cls.current_start_address + cls.current_rank_stride * i + cls.current_mem_offset - ) - if i != comm_rank: - # Fabric memory mapping - mem_handles[i] = _check_cu_result( - cuda.cuMemImportFromShareableHandle( - remote_handle_data, allocation_prop.requestedHandleTypes - ) - ) - - _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, mem_handles[i], 0)) - mapped_rank_ptrs.append(rank_ptr) - _check_cu_result(cuda.cuMemSetAccess(rank_ptr, aligned_size, [madesc], 1)) + mem_handles = cls._create_and_map_handles( + comm, + aligned_size, + cls.current_start_address, + cls.current_rank_stride, + cls.current_mem_offset, + ) except Exception: - # Clean up partially imported and mapped memory before propagating - # the failure to the caller. - for rank_ptr in reversed(mapped_rank_ptrs): - try: - _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) - except RuntimeError as e: - logger.warning("cuMemUnmap failed during error cleanup: %s", e) - for mem_handle in mem_handles: - if mem_handle is None: - continue - try: - _check_cu_result(cuda.cuMemRelease(mem_handle)) - except RuntimeError as e: - logger.warning("cuMemRelease failed during error cleanup: %s", e) if reserved_new_address: try: device_ptr = cuda.CUdeviceptr(cls.current_start_address) _check_cu_result( cuda.cuMemAddressFree(device_ptr, comm_size * cls.current_rank_stride) ) - except RuntimeError as e: - logger.warning("cuMemAddressFree failed during error cleanup: %s", e) + except RuntimeError as error: + logger.warning(f"cuMemAddressFree failed during error cleanup: {error}") else: ( cls.current_start_address, @@ -361,16 +597,18 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): cls.current_mem_offset, ) = previous_address_state raise - ptr = cls.current_start_address + cls.current_mem_offset stride = cls.current_rank_stride - cls.allocated_map[ptr] = ( - mapping, - aligned_size, - mem_handles, - cls.current_start_address, - cls.current_rank_stride, - cls.current_mem_offset, + cls.allocated_map[ptr] = _MnnvlAllocationRecord( + 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, + rank_stride=cls.current_rank_stride, + address_offset=cls.current_mem_offset, ) cls.address_refcnt[cls.current_start_address] = ( cls.address_refcnt.get(cls.current_start_address, 0) + 1 @@ -381,26 +619,197 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): @classmethod def close_mnnvl_memory(cls, ptr: int): - mapping, aligned_size, mem_handles, start_address, rank_stride, address_offset = ( - cls.allocated_map.pop(ptr) - ) - comm = cls.get_comm(mapping) - comm_size = comm.Get_size() - for i in range(comm_size): - rank_ptr = start_address + i * rank_stride + address_offset - _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) - _check_cu_result(cuda.cuMemRelease(mem_handles[i])) - cls.address_refcnt[start_address] -= 1 - - if cls.address_refcnt[start_address] == 0: - cls.address_refcnt.pop(start_address) - device_ptr = cuda.CUdeviceptr(start_address) - _check_cu_result(cuda.cuMemAddressFree(device_ptr, comm_size * rank_stride)) - if start_address == cls.current_start_address: + record = cls.allocated_map[ptr] + if record.state not in ( + _MnnvlAllocationState.MAPPED, + _MnnvlAllocationState.UNMAPPED, + ): + logger.warning( + f"Skipping cleanup of MNNVL allocation in terminal state {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 + + if cls.address_refcnt[record.start_address] == 0: + cls.address_refcnt.pop(record.start_address) + device_ptr = cuda.CUdeviceptr(record.start_address) + _check_cu_result( + cuda.cuMemAddressFree(device_ptr, record.comm_size * record.rank_stride) + ) + if record.start_address == cls.current_start_address: cls.current_start_address = 0 cls.current_rank_stride = 0 cls.current_mem_offset = 0 + @classmethod + def _unmap_and_release_handles(cls, record: _MnnvlAllocationRecord) -> None: + first_error = None + for rank in range(record.comm_size): + 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: + if first_error is None: + first_error = error + continue + try: + _check_cu_result(cuda.cuMemRelease(record.mem_handles[rank])) + except RuntimeError as error: + if first_error is None: + first_error = error + else: + record.mem_handles[rank] = None + if first_error is not None: + raise first_error + + def checkpoint_prepare(self) -> None: + """Detach local backing handles while retaining graph-visible VA. + + The engine checkpoint coordinator is responsible for quiescing and + invoking this operation on every rank. No collective follows the CUDA + mutation, so a local failure cannot strand peers in a trailing barrier. + """ + cls = type(self) + record = cls.allocated_map[self.ptr] + if record.state is _MnnvlAllocationState.UNMAPPED: + return + 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() + cls._unmap_and_release_handles(record) + record.mem_handles = [None] * record.comm_size + except Exception: + record.state = _MnnvlAllocationState.BROKEN + raise + record.state = _MnnvlAllocationState.UNMAPPED + + def checkpoint_fail_closed(self) -> None: + """Make a timed-out checkpoint allocation terminal. + + A timeout means ranks may have reached different checkpoint phases. + The allocation must not be reused by a later checkpoint attempt. + """ + record = type(self).allocated_map[self.ptr] + record.state = _MnnvlAllocationState.BROKEN + record.pending_comm = None + + def checkpoint_restore(self, comm: MnnvlCheckpointCommunicator) -> bool: + """Remap fresh handles while keeping data-path access disabled.""" + cls = type(self) + record = cls.allocated_map[self.ptr] + 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: + raise RuntimeError( + "Cannot restore MNNVL memory with a communicator that differs from " + "the graph-visible allocation layout: " + f"rank/size {comm_rank}/{comm_size} != " + f"{record.comm_rank}/{record.comm_size}" + ) + try: + comm_membership = tuple( + int(rank) + for rank in _checkpoint_allgather( + comm, + self.mapping.rank, + operation="communicator membership", + ) + ) + except Exception: + self.checkpoint_fail_closed() + raise + 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 + local_error = None + 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 TimeoutError: + record.state = _MnnvlAllocationState.BROKEN + raise + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + + try: + restore_errors = _checkpoint_allgather( + comm, + local_error, + operation="mapping readiness", + ) + except Exception: + self._checkpoint_restore_failed() + raise + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(restore_errors) + if error is not None + ] + if failed_ranks: + self._checkpoint_restore_failed() + raise RuntimeError( + "MNNVL checkpoint restore failed on one or more ranks:\n" + "\n".join(failed_ranks) + ) + 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): + if handle is None: + continue + 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( + f"Failed to unmap unpublished MNNVL restore for rank {rank}: {error}" + ) + try: + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning( + "Failed to release unpublished MNNVL restore handle " + f"for rank {rank}: {error}" + ) + else: + record.mem_handles[rank] = None + record.state = _MnnvlAllocationState.BROKEN + record.pending_comm = None + @staticmethod @functools.cache def support_nvlink(dev_id: int, need_all_up: bool = True): @@ -596,6 +1005,68 @@ def get_moe_prepare_workspace(mapping: Mapping): ) return MnnvlMoe.moe_prepare_workspace_tensor + @staticmethod + def checkpoint_prepare() -> None: + """Detach TRT-native two-sided MoE workspaces for checkpointing.""" + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace): + if workspace is not None: + workspace.checkpoint_prepare() + + @staticmethod + def checkpoint_restore(comm: MnnvlCheckpointCommunicator) -> None: + """Restore TRT-native two-sided MoE workspaces at their original virtual addresses.""" + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + 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 + restored_main_workspace = any( + workspace is MnnvlMoe.moe_workspace for workspace in restored_workspaces + ) + local_error = None + try: + if restored_main_workspace and 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() + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + readiness_errors = _checkpoint_allgather( + comm, + local_error, + operation="two-sided frontend readiness", + ) + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(readiness_errors) + if error is not None + ] + if failed_ranks: + raise RuntimeError( + "Native two-sided MoE restore failed on one or more ranks:\n" + + "\n".join(failed_ranks) + ) + 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: + """Reject kernel access while either native MoE workspace is detached.""" + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace): + if workspace is not None and not workspace.mapped: + raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") + @staticmethod def compute_target_rank_id( token_selected_experts: torch.Tensor, expert_count: int, ep_size: int diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 8d554921ef37..12e80d4591e3 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -28,12 +28,15 @@ import torch -from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlMemory +from tensorrt_llm._mnnvl_utils import (CftMnnvlMemory, + MnnvlCheckpointCommunicator, MnnvlMemory) from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, 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 @@ -370,48 +373,67 @@ def __init__( assert workspace_entry[ "can_use_cft_counted_writes"] == self.can_use_cft_counted_writes, "reuse workspace with different CFT mode" + workspace_state = workspace_entry self.mnnvl_mem = workspace_entry["mnnvl_mem"] self.workspace = workspace_entry["workspace"] - self.metainfo = workspace_entry["metainfo"] # 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 = workspace_entry self._workspace_state = workspace_state - 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._destroyed = False + self._workspace_registered = False + 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._workspace_registered = True + + @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 checkpoint_resource_key(self) -> int: + """Identify wrappers sharing the same MNNVL workspace lifecycle.""" + return id(self._workspace_lifecycle) 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 and getattr(self, "_workspace_registered", + False): + lifecycle.unregister(self) + self._workspace_registered = False def __del__(self) -> None: if not sys.is_finalizing(): @@ -446,6 +468,55 @@ def cft_initialize(self) -> None: f"CFT LE initialized (workspace-bound): ep_rank={self.ep_rank}, ep_size={self.ep_size}" ) + def _require_mapped(self) -> None: + if not self.mnnvl_mem.mapped: + raise RuntimeError( + "Native MoE All-to-All workspace handles are unmapped") + + def checkpoint_prepare(self) -> None: + """Collectively detach handles after every shared owner is idle.""" + if self.can_use_cft_counted_writes: + raise RuntimeError( + "Checkpointing a CFT-backed MoE All-to-All workspace is not supported" + ) + self._workspace_lifecycle.checkpoint_prepare() + + def checkpoint_restore( + self, + comm: MnnvlCheckpointCommunicator | None = None, + ) -> None: + """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. + """ + if comm is None: + comm = self.mnnvl_mem.comm + if comm is None: + raise RuntimeError( + "MNNVL workspace communicator is not initialized") + 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, + self.can_use_cft_counted_writes, + ), + ) + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self._state.phase == "idle" + + def _mnnvl_checkpoint_reset(self) -> None: + self.reset_state() + def dispatch(self, token_selected_experts: torch.Tensor, input_payloads: list[torch.Tensor], @@ -472,6 +543,7 @@ def dispatch(self, Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] """ + self._require_mapped() assert self._state.phase == "idle", "dispatch called twice without an intervening combine" reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" @@ -570,6 +642,7 @@ def combine( Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results """ + self._require_mapped() assert self._state.phase == "dispatched", "combine called before a successful dispatch" reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" @@ -617,6 +690,7 @@ def get_combine_payload_tensor_in_workspace( Return the combine payload tensor in the workspace, which could be used as the output of MoE kernel to avoid extra copy. Passing the returned tensor to combine lets the C++ op detect workspace ownership. """ + self._require_mapped() if self._state.phase != "dispatched": raise RuntimeError( "get_combine_payload_tensor_in_workspace called before a successful dispatch" diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py new file mode 100644 index 000000000000..07dde82ee935 --- /dev/null +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -0,0 +1,375 @@ +# 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 ( + MnnvlCheckpointCommunicator, + MnnvlMemory, + _checkpoint_allgather, +) +from tensorrt_llm._torch.alltoall_watchdog import ( + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + EPGroupHealthLike, +) + +_WORKSPACE_LIFECYCLE_KEY = "mnnvl_alltoall_workspace_lifecycle" + + +def _collect_active_ranks( + comm: MnnvlCheckpointCommunicator, + *, + local_clients_idle: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks whose local workspace clients are active.""" + clients_idle_by_rank = _checkpoint_allgather( + comm, + local_clients_idle, + operation="workspace idle readiness", + ) + 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: MnnvlCheckpointCommunicator, + *, + local_ready: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks that could not finish local restore work.""" + ready_by_rank = _checkpoint_allgather( + comm, + local_ready, + operation="frontend restore readiness", + ) + 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. + + ``PyExecutor`` invokes this resource hook from the engine sleep/wakeup + PREPARE/COMMIT/ABORT control path. That coordinator stops admission, drains + in-flight work, aggregates bounded per-rank results, and reopens admission + only after every rank has completed COMMIT. The local client and subgroup + idle checks here remain defense-in-depth validation while the executor is + parked. + """ + + 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(MnnvlCheckpointCommunicator | None, self._memory.comm) + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + try: + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self._ep_size, + ) + except TimeoutError: + self._memory.checkpoint_fail_closed() + self._stop_watchdog() + raise + 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: MnnvlCheckpointCommunicator, + 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/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index b03b26a9bb49..6930810ee2da 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -164,6 +164,16 @@ Still on old path (standalone, with embedded communication): Communication strategies are auto-selected at runtime by `CommunicationFactory` based on hardware and configuration. Skipped for `FUSED_COMM` backends. See `communication_factory.py` for selection logic and `base.py` for the `Communication` ABC. +Communication strategies whose native workspace must participate in executor +sleep/wakeup implement the runtime-checkable `CheckpointableCommunication` +protocol from `communication/base.py`. The public +`checkpoint_resource_key()` identifies wrappers that share one underlying +workspace so the executor invokes `checkpoint_prepare()` and +`checkpoint_restore()` once per resource. Discovery must use this protocol; +do not add concrete strategy checks or inspect private workspace attributes. +The executor request queue owns the persistent admission state around the +checkpoint operation and reopens admission only after a complete wakeup. + ### MegaMoE (`fused_moe/mega_moe/`) | File | Role | @@ -428,7 +438,7 @@ When adding new components, use these reference implementations: | New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | `capabilities`, `can_implement`, `run_moe`, `create_weights`, `load_weights`; then add the class to `moe_resolution.BACKEND_CANDIDATES`. Add a fixed `descriptor.identity` only for a one-implementation leaf class | | New `FUSED_COMM` Backend | `mega_moe/mega_moe_deepgemm.py` (`MegaMoEDeepGemm`), `mega_moe/mega_moe_cute_dsl.py` (`MegaMoECuteDsl`) | Same as above + override `scheduler_kind = MoESchedulerKind.FUSED_COMM` and `validate_configurable_moe` for backend-specific constraints. For NVFP4 CuteDSL specifically, mirror the `MegaMoECuteDsl` pattern: capability probe for the CUDA 13 Cutlass DSL runtime, JSON-friendly tactic dict, lazy kernel import via `cute_dsl_kernels/mega_moe_nvfp4/import_kernel()`, and `quantize_input` that short-circuits zero-token input. | | New Quantization Method | `quantization.py` → `FP8QDQFusedMoEMethod` | Subclass `FusedMoEMethod`, implement quant/dequant ops | -| New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine` | +| New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine`; also implement `CheckpointableCommunication` when native workspace mappings must follow executor sleep/wakeup | | New Scheduler | `moe_scheduler.py` (`ExternalCommMoEScheduler` / `FusedCommMoEScheduler`) | Subclass `MoEScheduler`, implement `forward`; add new `MoESchedulerKind` value and wire into `create_moe_scheduler` factory | | Backend Tests | `test_moe_backend.py` | Follow existing parametrize patterns | | Integration Tests | `test_moe_module.py` | Test Backend × Communication × EPLB combinations | diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py index 9858693c5f37..6e8780f35115 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py @@ -30,7 +30,7 @@ """ from .allgather_reducescatter import AllGatherReduceScatter -from .base import Communication +from .base import CheckpointableCommunication, Communication from .communication_factory import CommunicationFactory from .deep_ep import DeepEP from .deep_ep_low_latency import DeepEPLowLatency @@ -40,6 +40,7 @@ __all__ = [ # Base classes and types + "CheckpointableCommunication", "Communication", # Communication strategies "AllGatherReduceScatter", diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py b/tensorrt_llm/_torch/modules/fused_moe/communication/base.py index 0cfe8d5e5d7b..688e6849a564 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/base.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"); @@ -24,13 +24,35 @@ """ from abc import ABC, abstractmethod -from typing import List, Optional, Tuple +from collections.abc import Hashable +from typing import List, Optional, Protocol, Tuple, runtime_checkable import torch from tensorrt_llm.mapping import Mapping +@runtime_checkable +class CheckpointableCommunication(Protocol): + """MoE communication resource participating in engine sleep/wakeup. + + Implementations own their local checkpoint details and expose a stable key + so the executor invokes each shared workspace exactly once. + """ + + def checkpoint_resource_key(self) -> Hashable: + """Return the identity shared by wrappers using one checkpoint resource.""" + ... + + def checkpoint_prepare(self) -> None: + """Detach checkpoint-backed resources after global quiescence.""" + ... + + def checkpoint_restore(self) -> None: + """Restore checkpoint-backed resources before admission reopens.""" + ... + + class Communication(ABC): """ Abstract base class for MoE communication methods 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 0afe4697cf25..c7628cc1c7cf 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -127,6 +127,12 @@ def create_strategy( """ # Extract parameters from model_config mapping = model_config.mapping + if mapping.has_cp_helix(): + raise ValueError( + "MoE communication requires Helix context-parallel ranks to be " + "repurposed through Mapping.repurpose_helix_cp_to_tp() before " + "strategy selection" + ) if hidden_size is None: hidden_size = model_config.pretrained_config.hidden_size act_dtype = model_config.torch_dtype 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 8bc2754a58c1..be6e75c40dc2 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 +import sys from typing import Callable, Dict, List, Optional, Tuple import torch -from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlMemory +from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlCheckpointCommunicator, MnnvlMemory from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, @@ -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 @@ -296,6 +298,11 @@ def __init__( alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkOneSided does not support Helix context parallelism because " + "its MNNVL communicator covers only the tensor-parallel group" + ) if self.mapping.world_size != self.ep_size: raise RuntimeError("Currently NVLinkOneSided only supports pure EP for MoE.") @@ -396,7 +403,7 @@ def __init__( workspace_state = NVLinkOneSided._WORKSPACES.get(self._workspace_key) memory_cls = CftMnnvlMemory if self.can_use_cft_counted_writes else MnnvlMemory - + 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." @@ -429,7 +436,6 @@ def __init__( "metainfo": metainfo, "cft_initialized": False, } - NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: expected_workspace_state = { "workspace_size_per_rank": self.workspace_size_per_rank, @@ -449,41 +455,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 CFT Logical Endpoints by binding the LE to the workspace. # The LE IS the workspace — no separate allocation or payload layout needed. @@ -503,6 +521,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: """ @@ -517,17 +553,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 and getattr(self, "_workspace_registered", False): + 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() @@ -545,10 +583,28 @@ 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 + # Finalizers cannot safely perform collective cache eviction because + # their order is nondeterministic across ranks. + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None and getattr(self, "_workspace_registered", False): + 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. @@ -574,6 +630,57 @@ def use_cft_for_combine(self, runtime_max_tokens_per_rank: int) -> bool: runtime_max_tokens_per_rank, ) + def _require_mapped(self) -> None: + if not self.mnnvl_mem.mapped: + raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") + + def checkpoint_resource_key(self) -> int: + """Identify wrappers sharing the same MNNVL workspace lifecycle.""" + return id(self._require_workspace_lifecycle()) + + def checkpoint_prepare(self) -> None: + """Collectively detach handles after every shared owner is idle.""" + if self.can_use_cft_counted_writes: + raise RuntimeError( + "Checkpointing a CFT-backed MoE All-to-All workspace is not supported" + ) + self._require_workspace_lifecycle().checkpoint_prepare() + + def checkpoint_restore( + self, + comm: MnnvlCheckpointCommunicator | None = None, + ) -> None: + """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. + """ + if comm is None: + comm = self.mnnvl_mem.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + 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, + self.can_use_cft_counted_writes, + ), + ) + + 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( self, hidden_states: torch.Tensor, @@ -604,6 +711,7 @@ def dispatch( Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) Each tensor has shape [ep_size, max_tokens_per_rank, ...] """ + self._require_mapped() if self._dispatch_state.get("phase") == "dispatched": raise RuntimeError("dispatch called twice without an intervening combine") reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) @@ -744,6 +852,7 @@ def combine( Combined output tensor [local_num_tokens, hidden_size] """ + self._require_mapped() if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError("combine called before a successful dispatch") reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) @@ -840,6 +949,7 @@ def get_combine_payload_tensor_in_workspace( Returns: Tensor view into workspace [ep_size, max_tokens_per_rank, hidden_size] """ + self._require_mapped() if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError( "get_combine_payload_tensor_in_workspace called before a successful 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 61d03b3a973b..2582a965490f 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 @@ -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"); @@ -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._mnnvl_utils import MnnvlCheckpointCommunicator, 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, @@ -52,6 +56,11 @@ def __init__( alltoall_result_do_sum: bool = False, ): super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkTwoSided does not support Helix context parallelism because " + "its MNNVL communicator covers only the tensor-parallel group" + ) # Store needed parameters self.num_experts = num_experts @@ -76,6 +85,7 @@ def __init__( # Initialize dispatch state self._dispatch_state = {} + self._INSTANCES.add(self) @staticmethod def is_platform_supported() -> bool: @@ -99,6 +109,67 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + def checkpoint_resource_key(self) -> int: + """Identify the process-global TRT-native two-sided workspaces.""" + return id(MnnvlMoe) + + def checkpoint_prepare(self) -> None: + """Detach TRT-native two-sided workspaces after global quiescence.""" + 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") + try: + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self.ep_size, + ) + except TimeoutError: + for candidate in workspaces: + if candidate is not None: + candidate.checkpoint_fail_closed() + raise + 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: MnnvlCheckpointCommunicator | None = None, + ) -> None: + """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. + """ + workspace = MnnvlMoe.moe_workspace or MnnvlMoe.moe_prepare_workspace + if comm is None and workspace is not None: + comm = workspace.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + 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) + if not restore_required: + return + for instance in self._INSTANCES: + instance._dispatch_state = {} + def prepare_dispatch( self, token_selected_slots: torch.Tensor, @@ -108,6 +179,7 @@ def prepare_dispatch( """ NVLINK two-sided comm prepare dispatch: gather EPLB statistics and prepare alltoall_info. """ + MnnvlMoe.require_mapped() all_rank_max_num_tokens = max(all_rank_num_tokens) top_k = token_selected_slots.shape[1] @@ -144,6 +216,7 @@ def dispatch( """ NVLINK two-sided comm dispatch (post-quant, uses alltoall_info from prepare_dispatch). """ + MnnvlMoe.require_mapped() # Read alltoall_info from dispatch_state (set by prepare_dispatch) alltoall_info = self._dispatch_state.get("alltoall_info") if alltoall_info is None: @@ -189,6 +262,7 @@ def combine( """ NVLINK two-sided comm combine - reads from self._dispatch_state. """ + MnnvlMoe.require_mapped() if isinstance(final_hidden_states, list): final_hidden_states = final_hidden_states[0] @@ -204,4 +278,5 @@ def combine( do_reduce=self.alltoall_result_do_sum, ) + self._dispatch_state = {} return final_hidden_states diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py index 6963a046dd68..b3aaa2f8ad0e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py @@ -84,6 +84,11 @@ def __init__( alltoall_result_do_sum: bool = False, ): super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkTwoSidedFlashinfer does not support Helix context parallelism " + "because its MNNVL communicator covers only the tensor-parallel group" + ) # Store needed parameters self.num_experts = num_experts diff --git a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py index 8dcb325c2775..7d1b860c5cd4 100644 --- a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py +++ b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py @@ -1,5 +1,6 @@ import dataclasses import datetime +import enum import queue import threading import time @@ -15,6 +16,16 @@ CONTROL_REQUEST_ID = -2 +class RequestAdmissionState(enum.Enum): + """Persistent request-admission state across engine sleep and wakeup.""" + + RUNNING = "running" + PARKING = "parking" + PARKED = "parked" + WAKING = "waking" + FAILED = "failed" + + @dataclasses.dataclass class RequestQueueItem: id: int @@ -61,8 +72,97 @@ def __init__( self.enable_iter_perf_stats = enable_iter_perf_stats self.start_times = {} self.active = True + self.admission_state = RequestAdmissionState.RUNNING + self._pending_sleep_tags: frozenset[str] = frozenset() + self._parked_tags: frozenset[str] = frozenset() + self._pending_wakeup_tags: frozenset[str] = frozenset() self.batch_wait_timeout_ms = batch_wait_timeout_ms + def _transition_admission_state( + self, + expected: RequestAdmissionState, + target: RequestAdmissionState, + ) -> None: + if self.admission_state is not expected: + raise RuntimeError( + "Invalid executor admission transition: " + f"{self.admission_state.value} -> {target.value}; " + f"expected {expected.value}") + self.admission_state = target + + def begin_sleep_transition(self, tags: Iterable[str]) -> None: + """Close admission before the executor starts draining for sleep.""" + with self.enqueue_lock: + sleep_tags = frozenset(tags) + self._transition_admission_state( + RequestAdmissionState.RUNNING, + RequestAdmissionState.PARKING, + ) + self._pending_sleep_tags = sleep_tags + + def complete_sleep_transition(self) -> None: + """Publish PARKED only after every rank has completed sleep.""" + with self.enqueue_lock: + self._transition_admission_state( + RequestAdmissionState.PARKING, + RequestAdmissionState.PARKED, + ) + self._parked_tags = self._pending_sleep_tags + self._pending_sleep_tags = frozenset() + + def abort_sleep_transition(self) -> None: + """Reopen admission after a recoverable pre-mutation sleep failure.""" + with self.enqueue_lock: + if self.admission_state is RequestAdmissionState.FAILED: + return + self._transition_admission_state( + RequestAdmissionState.PARKING, + RequestAdmissionState.RUNNING, + ) + self._pending_sleep_tags = frozenset() + + def begin_wakeup_transition(self, tags: Iterable[str]) -> None: + """Keep admission closed while parked resources are restored.""" + with self.enqueue_lock: + wakeup_tags = frozenset(tags) + self._transition_admission_state( + RequestAdmissionState.PARKED, + RequestAdmissionState.WAKING, + ) + self._pending_wakeup_tags = wakeup_tags + + def complete_wakeup_transition(self) -> None: + """Reopen admission only after all originally parked tags are restored.""" + with self.enqueue_lock: + remaining_tags = self._parked_tags - self._pending_wakeup_tags + target = RequestAdmissionState.PARKED + if not remaining_tags: + target = RequestAdmissionState.RUNNING + self._transition_admission_state(RequestAdmissionState.WAKING, + target) + self._parked_tags = remaining_tags + self._pending_wakeup_tags = frozenset() + + def abort_wakeup_transition(self) -> None: + """Return to PARKED after a recoverable pre-mutation wakeup failure.""" + with self.enqueue_lock: + if self.admission_state is RequestAdmissionState.FAILED: + return + self._transition_admission_state( + RequestAdmissionState.WAKING, + RequestAdmissionState.PARKED, + ) + self._pending_wakeup_tags = frozenset() + + def fail_sleep_wakeup_transition(self) -> None: + """Permanently close admission after sleep/wakeup may have mutated state.""" + with self.enqueue_lock: + self.admission_state = RequestAdmissionState.FAILED + + def get_admission_state(self) -> RequestAdmissionState: + with self.enqueue_lock: + return self.admission_state + def _get_request_id(self, request: Optional[ExecutorRequest] = None): # if request has a disagg_request_id, use it as request id so that # corresponding context and generation requests have the same request id @@ -93,6 +193,10 @@ def _enqueue_impl(self, requests: Iterable[ExecutorRequest]) -> List[int]: req_ids = [] with self.enqueue_lock: assert self.active, "PyExecutor has already been shutdown." + if self.admission_state is not RequestAdmissionState.RUNNING: + raise RuntimeError( + "Cannot enqueue requests while executor admission is " + f"{self.admission_state.value}") start_time = time.time() for request in requests: req_id = self._get_request_id(request) @@ -140,6 +244,12 @@ def enqueue_shutdown_request(self): self.active = False def can_enqueue_request(self) -> bool: + with self.enqueue_lock: + return (self.active and self.dist.rank == 0 + and self.admission_state is RequestAdmissionState.RUNNING) + + def can_enqueue_control_request(self) -> bool: + """Return whether rank zero can enqueue shutdown/control sentinels.""" with self.enqueue_lock: return self.active and self.dist.rank == 0 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..c50faf9b0776 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,7 +42,8 @@ from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata -from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy +from tensorrt_llm.llmapi.llm_args import (ExecutorMemoryType, PeftCacheConfig, + WaitingQueuePolicy) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError @@ -68,7 +69,8 @@ from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .error_classification import ErrorBudget -from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem +from .executor_request_queue import (ExecutorRequestQueue, + RequestAdmissionState, RequestQueueItem) from .guided_decoder import GuidedDecoder from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits @@ -104,6 +106,9 @@ if TYPE_CHECKING: from ray.actor import ActorHandle + from ..modules.fused_moe.communication.base import \ + CheckpointableCommunication + _UNBOUNDED_STATS_MAX_LEN = -1 @@ -233,23 +238,16 @@ def _recv_sleep_wakeup_ack_until(comm, if (expected_op_id is not None and ack.get("op_id") != expected_op_id): logger.warning( - "Ignoring stale sleep/wakeup ACK from rank %d for op_id=%s " - "(expected op_id=%s).", - source, - ack.get("op_id"), - expected_op_id, + f"Ignoring stale sleep/wakeup ACK from rank {source} for " + f"op_id={ack.get('op_id')} (expected op_id={expected_op_id})." ) continue ack_phase = ack.get("phase") if (expected_phase is not None and ack_phase is not None and ack_phase != expected_phase): logger.warning( - "Ignoring stale sleep/wakeup ACK from rank %d for phase=%s " - "(expected phase=%s).", - source, - ack_phase, - expected_phase, - ) + f"Ignoring stale sleep/wakeup ACK from rank {source} for " + f"phase={ack_phase} (expected phase={expected_phase}).") continue return ack if time.monotonic() >= deadline: @@ -1343,10 +1341,8 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: return if self.dist.rank == 0: - logger.info( - "Sending shutdown to %d sleep/wakeup listener thread(s).", - self.dist.world_size - 1, - ) + logger.info(f"Sending shutdown to {self.dist.world_size - 1} " + "sleep/wakeup listener thread(s).") shutdown_ack_deadline = (time.monotonic() + _SLEEP_WAKEUP_ACK_TIMEOUT_S) shutdown_errors = [] @@ -1367,7 +1363,7 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: f"rank {dest} shutdown send failed: {exc}") logger.warning( "Failed to send sleep/wakeup listener shutdown to " - "rank %d: %s", dest, exc) + f"rank {dest}: {exc}") for src in shutdown_ranks: try: ack = _recv_sleep_wakeup_ack_until(self._sleep_wakeup_comm, @@ -1378,7 +1374,7 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: f"rank {src} shutdown ACK recv failed: {exc}") logger.warning( "Failed to receive sleep/wakeup listener shutdown ACK " - "from rank %d: %s", src, exc) + f"from rank {src}: {exc}") continue if ack.get("status") != "ok": shutdown_errors.append( @@ -1386,18 +1382,16 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: or f"rank {src} returned unknown shutdown ACK") if shutdown_errors: logger.warning( - "Sleep/wakeup listener shutdown completed with errors: %s", - "; ".join(shutdown_errors)) + "Sleep/wakeup listener shutdown completed with errors: " + f"{'; '.join(shutdown_errors)}") elif self._sleep_wakeup_listener_thread is not None: self._sleep_wakeup_listener_thread.join( timeout=_SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S) if self._sleep_wakeup_listener_thread.is_alive(): logger.warning( - "Sleep/wakeup listener thread did not exit within %.1f " - "seconds on rank %d.", - _SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S, - self.dist.rank, - ) + "Sleep/wakeup listener thread did not exit within " + f"{_SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S:.1f} seconds on " + f"rank {self.dist.rank}.") def _record_sleep_wakeup_abort(self, control_id: str, error_msg: str) -> None: @@ -1545,6 +1539,35 @@ def can_enqueue_requests(self) -> bool: """ return self.executor_request_queue.can_enqueue_request() + def begin_sleep_transition(self, tags: list[ExecutorMemoryType]) -> None: + self.executor_request_queue.begin_sleep_transition(tag.value + for tag in tags) + + def complete_sleep_transition(self) -> None: + self.executor_request_queue.complete_sleep_transition() + + def abort_sleep_transition(self) -> None: + self.executor_request_queue.abort_sleep_transition() + + def begin_wakeup_transition(self, tags: list[ExecutorMemoryType]) -> None: + self.executor_request_queue.begin_wakeup_transition(tag.value + for tag in tags) + + def complete_wakeup_transition(self) -> None: + self.executor_request_queue.complete_wakeup_transition() + + def abort_wakeup_transition(self) -> None: + self.executor_request_queue.abort_wakeup_transition() + + def fail_sleep_wakeup_transition(self) -> None: + self.executor_request_queue.fail_sleep_wakeup_transition() + + def get_request_admission_state(self) -> RequestAdmissionState: + return self.executor_request_queue.get_admission_state() + + def can_shutdown(self) -> bool: + return self.executor_request_queue.can_enqueue_control_request() + def get_latest_iteration_stats(self): """ Returns the per-iterations statistics computed since last call to this method. @@ -3023,8 +3046,8 @@ def _sleep_wakeup_listener_loop(self) -> None: tag=_SleepWakeupTag.ACTION) if msg.get("action") == _SleepWakeupAction.SHUTDOWN: logger.debug( - "Sleep/wakeup listener (rank %d): received shutdown, " - "exiting.", self.dist.rank) + f"Sleep/wakeup listener (rank {self.dist.rank}): " + "received shutdown, exiting.") self._sleep_wakeup_comm.send( { "status": "ok", @@ -3040,7 +3063,9 @@ def _sleep_wakeup_listener_loop(self) -> None: action = msg.get("action", "") op_id = msg.get("op_id") error_msg = None + abort_acknowledged = False release_control_request = True + has_mnnvl_resources = False try: # Decode tags inside the try so KeyError / ValueError from # a malformed message still results in an error ACK being @@ -3068,11 +3093,14 @@ def _sleep_wakeup_listener_loop(self) -> None: None) if op_id is None: release_control_request = False - elif (not self.control_request_barrier.is_set() - or active_control_id != op_id): - self._record_sleep_wakeup_abort(op_id, error_msg) - release_control_request = False - logger.warning("Sleep/wakeup listener: %s", error_msg) + else: + if (not self.control_request_barrier.is_set() + or active_control_id != op_id): + self._record_sleep_wakeup_abort( + op_id, error_msg) + release_control_request = False + abort_acknowledged = True + logger.warning(f"Sleep/wakeup listener: {error_msg}") elif action in (_SleepWakeupAction.PREPARE, _SleepWakeupAction.COMMIT, _SleepWakeupAction.SLEEP, @@ -3085,15 +3113,19 @@ def _sleep_wakeup_listener_loop(self) -> None: f"stale control message for op_id={op_id}; " f"active control_id={active_control_id}") release_control_request = False - logger.warning("Sleep/wakeup listener: %s", - error_msg) + logger.warning( + f"Sleep/wakeup listener: {error_msg}") else: torch.cuda.synchronize() if action == _SleepWakeupAction.PREPARE: # Prepared means this rank is quiesced and # ready to commit, but VMM state is unchanged. + has_mnnvl_resources = ( + self._has_mnnvl_checkpoint_resources(tags)) release_control_request = False elif target_action == _SleepWakeupAction.SLEEP: + self._run_mnnvl_checkpoint_resources( + target_action, tags) release_with_tag(*tags) torch.cuda.synchronize() gc.collect() @@ -3101,23 +3133,23 @@ def _sleep_wakeup_listener_loop(self) -> None: elif target_action == _SleepWakeupAction.WAKEUP: materialize_with_tag(*tags) torch.cuda.synchronize() + self._run_mnnvl_checkpoint_resources( + target_action, tags) else: error_msg = ( f"unknown target action '{target_action}'") logger.warning( - "Sleep/wakeup listener: %s, ignoring.", - error_msg) + f"Sleep/wakeup listener: {error_msg}, ignoring." + ) else: error_msg = f"unknown action '{action}'" - logger.warning("Sleep/wakeup listener: %s, ignoring.", - error_msg) + logger.warning( + f"Sleep/wakeup listener: {error_msg}, ignoring.") except (KeyError, TypeError, ValueError, RuntimeError, - torch.OutOfMemoryError) as exc: + TimeoutError, torch.OutOfMemoryError) as exc: error_msg = (f"rank {self.dist.rank} '{action}' failed: " f"{exc}\n{traceback.format_exc()}") - logger.error("Sleep/wakeup listener: error executing '%s':", - action, - exc_info=True) + logger.error(f"Sleep/wakeup listener: {error_msg}") finally: # Always ACK so rank-0 does not deadlock; carry error # details so rank-0 can raise after all ranks respond. @@ -3129,14 +3161,9 @@ def _sleep_wakeup_listener_loop(self) -> None: exc = sys.exc_info()[1] error_msg = ( f"rank {self.dist.rank} '{action}' failed with " - f"uncaught {type(exc).__name__}: {exc!r}") - logger.error( - "Sleep/wakeup listener: uncaught exception on " - "rank %d during '%s':", - self.dist.rank, - action, - exc_info=True, - ) + f"uncaught {type(exc).__name__}: {exc!r}\n" + f"{traceback.format_exc()}") + logger.error(f"Sleep/wakeup listener: {error_msg}") # Unblock the executor loop that is waiting in # _handle_control_request(). Clear control_request_barrier # first so that it is clean for the next sleep/wakeup cycle @@ -3150,19 +3177,78 @@ def _sleep_wakeup_listener_loop(self) -> None: # exiting control_action() and broadcasting new requests # before our executor loop has cleared its control barrier # and is ready to participate in the next collective. + ack_error = None if abort_acknowledged else error_msg + ack = { + "status": "ok" if ack_error is None else "error", + "error": ack_error, + "op_id": op_id, + "phase": action, + "has_mnnvl_resources": has_mnnvl_resources, + } + if abort_acknowledged: + ack["reason"] = error_msg self._sleep_wakeup_comm.send( - { - "status": "ok" if error_msg is None else "error", - "error": error_msg, - "op_id": op_id, - "phase": action, - }, + ack, dest=0, tag=_SleepWakeupTag.ACK, ) finally: set_thread_local_mpi_comm(None) + def _mnnvl_checkpoint_resources( + self, + tags: list[ExecutorMemoryType], + ) -> list["CheckpointableCommunication"]: + """Return unique native MNNVL MoE resources selected by engine tags.""" + from tensorrt_llm._torch.modules.fused_moe.communication.base import \ + CheckpointableCommunication + + selected_engines = [] + if ExecutorMemoryType.MODEL_ENGINE_MAIN in tags: + selected_engines.append(self.model_engine) + if ExecutorMemoryType.MODEL_ENGINE_DRAFT in tags and self.draft_model_engine is not None: + selected_engines.append(self.draft_model_engine) + + resources: list[CheckpointableCommunication] = [] + seen = set() + for engine in selected_engines: + model = getattr(engine, "model", None) + if model is None: + continue + for module in model.modules(): + resource = getattr(module, "comm", None) + if not isinstance(resource, CheckpointableCommunication): + continue + key = resource.checkpoint_resource_key() + if key in seen: + continue + seen.add(key) + resources.append(resource) + return resources + + def _has_mnnvl_checkpoint_resources( + self, + tags: list[ExecutorMemoryType], + ) -> bool: + return bool(self._mnnvl_checkpoint_resources(tags)) + + def _run_mnnvl_checkpoint_resources( + self, + action: _SleepWakeupAction, + tags: list[ExecutorMemoryType], + ) -> None: + """Execute process-local native MNNVL hooks while the engine is parked.""" + resources = self._mnnvl_checkpoint_resources(tags) + if action == _SleepWakeupAction.SLEEP: + for resource in resources: + resource.checkpoint_prepare() + return + if action == _SleepWakeupAction.WAKEUP: + for resource in resources: + resource.checkpoint_restore() + return + raise ValueError(f"unknown MNNVL checkpoint action '{action}'") + def _ring_broadcast_sample_state( self, executed_batch: Optional[BatchStatePP], @@ -4490,10 +4576,8 @@ def _handle_control_request(self): pending_abort = self._pop_sleep_wakeup_abort(control_id) if pending_abort is not None: logger.warning( - "[control_action] skipping aborted control request %s: %s", - control_id, - pending_abort, - ) + f"[control_action] skipping aborted control request {control_id}: " + f"{pending_abort}") self.control_request_barrier.set() self.control_request_barrier.clear() self._active_control_id = None diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 098d4674603d..bc63be7d2e56 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -22,7 +22,7 @@ import weakref from pathlib import Path from queue import Empty, Queue -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union import torch @@ -633,7 +633,7 @@ def _check_sleep_wakeup_preconditions(self, method: str) -> None: def _multi_rank_sleep_wakeup( self, - action: str, + action: Literal["sleep", "wakeup"], tags: list[ExecutorMemoryType], ) -> None: """Coordinate a sleep or wakeup operation across all MPI ranks. @@ -647,12 +647,17 @@ def _multi_rank_sleep_wakeup( 2. Send PREPARE to every non-rank-0 rank via the dedicated ``_sleep_wakeup_comm`` communicator. Peers quiesce and ACK without changing VMM state. - 3. Execute the VMM operation (``release_with_tag`` or - ``materialize_with_tag``) locally on rank-0. - 4. Send COMMIT to prepared peers and collect ACKs after their local VMM - operations. If PREPARE or rank-0 local execution fails, send ABORT - so peers leave the control barrier without changing VMM state. - 5. Exit ``control_action()``, resuming rank-0's event loop. + 3. When native MNNVL MoE resources are selected, send COMMIT to every + prepared peer before rank-0 enters the subgroup checkpoint + collectives. Otherwise execute the local VMM operation first. + 4. Execute the selected VMM and MNNVL operations and collect bounded + COMMIT ACKs. If PREPARE or the initial rank-0 synchronization fails, + send ABORT before any rank changes VMM state. Once COMMIT or local + mutation starts, any error fail-stops the distributed worker because + rank state can no longer be reconciled safely. + 5. Exit ``control_action()`` only after a successful operation or a + recoverable pre-COMMIT abort. The public ``sleep()``/``wakeup()`` + wrapper owns the persistent admission transition around this helper. Args: action: ``"sleep"`` or ``"wakeup"``. @@ -683,6 +688,9 @@ def _multi_rank_sleep_wakeup( tag_strings = [t.value for t in tags] op_id = uuid.uuid4().hex target_action = _SleepWakeupAction(action) + has_mnnvl_resources = getattr(self.engine, + "_has_mnnvl_checkpoint_resources", + lambda _tags: False)(tags) prepare_msg = { "action": _SleepWakeupAction.PREPARE, "target_action": target_action, @@ -707,9 +715,14 @@ def _multi_rank_sleep_wakeup( errors = [] local_error = None abort_sent = False + commit_ranks = [] + commit_sent_early = False + local_commit_started = False + abort_incomplete = False def send_abort(reason: str, ranks: Optional[list[int]] = None) -> list[int]: + nonlocal abort_incomplete abort_ranks = [] abort_dests = ranks if ranks is not None else range( 1, world_size) @@ -728,18 +741,20 @@ def send_abort(reason: str, ) abort_ranks.append(abort_dest) except Exception as abort_exc: + abort_incomplete = True abort_error = ( "rank 0 failed to send sleep/wakeup abort " f"to rank {abort_dest}: {abort_exc}") errors.append(abort_error) logger.error( - "_multi_rank_sleep_wakeup: %s", - abort_error, - exc_info=True, - ) + f"_multi_rank_sleep_wakeup: {abort_error}\n" + f"{traceback.format_exc()}") return abort_ranks - def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: + def drain_acks(ranks: list[int], + phase: _SleepWakeupAction) -> list[dict]: + nonlocal abort_incomplete + received_acks = [] ack_deadline = time.monotonic() + _SLEEP_WAKEUP_ACK_TIMEOUT_S for src in ranks: try: @@ -749,21 +764,50 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: expected_op_id=op_id, expected_phase=phase) except Exception as exc: + if phase == _SleepWakeupAction.ABORT: + abort_incomplete = True errors.append( f"rank 0 failed to receive {phase} ACK from " f"rank {src}: {exc}") logger.error( - "_multi_rank_sleep_wakeup: failed to receive %s " - "ACK from rank %d", - phase, - src, - exc_info=True, + "_multi_rank_sleep_wakeup: failed to receive " + f"{phase} ACK from rank {src}\n{traceback.format_exc()}" ) continue + received_acks.append(ack) if ack.get("status") != "ok": + if phase == _SleepWakeupAction.ABORT: + abort_incomplete = True errors.append( ack.get("error") or f"rank {src} returned unknown {phase} ACK") + return received_acks + + def send_commits() -> list[int]: + sent_ranks = [] + failed_ranks = [] + for dest in prepared_ranks: + try: + sleep_wakeup_comm.send( + commit_msg, + dest=dest, + tag=_SleepWakeupTag.ACTION, + ) + sent_ranks.append(dest) + except Exception as exc: + commit_error = ( + f"rank 0 failed to send '{action}' commit to " + f"rank {dest}: {exc}") + errors.append(commit_error) + failed_ranks.append(dest) + logger.error( + f"_multi_rank_sleep_wakeup: {commit_error}\n" + f"{traceback.format_exc()}") + if failed_ranks: + abort_ranks = send_abort("\n".join(errors), + ranks=failed_ranks) + drain_acks(abort_ranks, _SleepWakeupAction.ABORT) + return sent_ranks try: # Phase 1: prepare peers. A prepared peer has reached the @@ -781,11 +825,8 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: f"rank 0 failed to send '{action}' prepare to rank " f"{dest}: {exc}") errors.append(send_error) - logger.error( - "_multi_rank_sleep_wakeup: %s", - send_error, - exc_info=True, - ) + logger.error(f"_multi_rank_sleep_wakeup: {send_error}\n" + f"{traceback.format_exc()}") abort_ranks = send_abort(send_error) abort_sent = True drain_acks(prepared_ranks, _SleepWakeupAction.PREPARE) @@ -793,15 +834,30 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: break if not errors: - drain_acks(prepared_ranks, _SleepWakeupAction.PREPARE) + prepare_acks = drain_acks(prepared_ranks, + _SleepWakeupAction.PREPARE) + has_mnnvl_resources = has_mnnvl_resources or any( + ack.get("has_mnnvl_resources", False) + for ack in prepare_acks) if not errors: - # Execute locally on rank-0. Only CUDA/VMM errors are - # captured as local_error. Peers are still prepared but - # uncommitted, so local failure can abort them without - # changing their VMM state. + # MNNVL resource hooks contain subgroup handle exchange, + # so peers must enter COMMIT before rank 0 executes them. + if has_mnnvl_resources: + # A failed MPI send has uncertain delivery, so any + # attempted COMMIT requires the bounded local phase. + commit_sent_early = bool(prepared_ranks) + commit_ranks = send_commits() + + if not errors or commit_sent_early: torch.cuda.synchronize() + local_commit_started = True if action == _SleepWakeupAction.SLEEP: + run_mnnvl = (getattr( + self.engine, "_run_mnnvl_checkpoint_resources", + None) if has_mnnvl_resources else None) + if run_mnnvl is not None: + run_mnnvl(target_action, tags) release_with_tag(*tags) torch.cuda.synchronize() gc.collect() @@ -809,53 +865,60 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: else: materialize_with_tag(*tags) torch.cuda.synchronize() - except (RuntimeError, torch.OutOfMemoryError) as exc: + run_mnnvl = (getattr( + self.engine, "_run_mnnvl_checkpoint_resources", + None) if has_mnnvl_resources else None) + if run_mnnvl is not None: + run_mnnvl(target_action, tags) + except Exception as exc: local_error = (f"rank 0 '{action}' failed: {exc}\n" f"{traceback.format_exc()}") - logger.error( - "_multi_rank_sleep_wakeup: rank-0 local %s failed:", - action, - exc_info=True, - ) + logger.error(f"_multi_rank_sleep_wakeup: {local_error}") finally: if local_error: errors.append(local_error) - if errors and prepared_ranks and not abort_sent: + if commit_sent_early: + drain_acks(commit_ranks, _SleepWakeupAction.COMMIT) + + if (errors and prepared_ranks and not abort_sent + and not commit_sent_early): abort_ranks = send_abort("\n".join(errors)) drain_acks(abort_ranks, _SleepWakeupAction.ABORT) - elif not errors: - commit_ranks = [] - commit_failed_ranks = [] - for dest in prepared_ranks: - try: - sleep_wakeup_comm.send( - commit_msg, - dest=dest, - tag=_SleepWakeupTag.ACTION, - ) - commit_ranks.append(dest) - except Exception as exc: - commit_error = ( - f"rank 0 failed to send '{action}' commit to " - f"rank {dest}: {exc}") - errors.append(commit_error) - commit_failed_ranks.append(dest) - logger.error( - "_multi_rank_sleep_wakeup: %s", - commit_error, - exc_info=True, - ) - if commit_failed_ranks: - abort_ranks = send_abort("\n".join(errors), - ranks=commit_failed_ranks) - drain_acks(abort_ranks, _SleepWakeupAction.ABORT) + elif not errors and not commit_sent_early: + commit_ranks = send_commits() drain_acks(commit_ranks, _SleepWakeupAction.COMMIT) if errors: - raise RuntimeError( + operation_error = RuntimeError( f"{action}() failed on {len(errors)} rank(s):\n" + "\n".join(errors)) + if commit_sent_early or local_commit_started or abort_incomplete: + self._fail_stop_divergent_sleep_wakeup(operation_error) + raise operation_error + + def _fail_stop_divergent_sleep_wakeup(self, error: RuntimeError) -> None: + """Stop every rank after a sleep/wakeup operation may have diverged.""" + from tensorrt_llm._torch.pyexecutor.hang_detector import \ + propagate_hard_kill + + self._set_fatal_error(error) + if self.engine is not None: + fail_transition = getattr( + self.engine, + "fail_sleep_wakeup_transition", + None, + ) + if fail_transition is not None: + fail_transition() + if getattr(self.engine, "_fatal_error", None) is None: + self.engine._fatal_error = error + self.engine.is_shutdown = True + try: + logger.critical("Distributed sleep/wakeup state may have diverged; " + f"hard-killing all ranks: {error}") + finally: + propagate_hard_kill() def sleep(self, sleep_tags: list[str]) -> None: """Release GPU virtual memory for the specified memory type tags. @@ -882,9 +945,9 @@ def sleep(self, sleep_tags: list[str]) -> None: value strings (e.g. ``["kv_cache"]``). Returns: - None. The call is synchronous; when it returns all requested - VMM-tagged allocations have been released on every rank and the - event loop has been resumed. + None. The call is synchronous; when it returns all requested + VMM-tagged allocations have been released on every rank and request + admission remains closed until ``wakeup()`` succeeds. Raises: ValueError: If the backend is not ``"pytorch"`` or @@ -896,15 +959,30 @@ def sleep(self, sleep_tags: list[str]) -> None: tags = [ExecutorMemoryType(tag) for tag in sleep_tags] logger.info(f"Sleep: {tags}") - if self.llm_args.parallel_config.world_size > 1: - self._multi_rank_sleep_wakeup("sleep", tags) - else: - with self.engine._sleep_wakeup_lock, self.engine.control_action(): - torch.cuda.synchronize() - release_with_tag(*tags) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() + self.engine.begin_sleep_transition(tags) + local_mutation_started = False + try: + if self.llm_args.parallel_config.world_size > 1: + self._multi_rank_sleep_wakeup("sleep", tags) + else: + with self.engine._sleep_wakeup_lock, self.engine.control_action( + ): + torch.cuda.synchronize() + local_mutation_started = True + release_with_tag(*tags) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + except Exception: + if local_mutation_started: + self.engine.fail_sleep_wakeup_transition() + self.engine.abort_sleep_transition() + raise + try: + self.engine.complete_sleep_transition() + except Exception: + self.engine.fail_sleep_wakeup_transition() + raise def wakeup(self, wakeup_tags: list[str]) -> None: """Materialize GPU virtual memory for the specified memory type tags. @@ -920,9 +998,10 @@ def wakeup(self, wakeup_tags: list[str]) -> None: value strings (e.g. ``["kv_cache"]``). Returns: - None. The call is synchronous; when it returns all requested - VMM-tagged allocations have been materialized on every rank and the - event loop has been resumed. + None. The call is synchronous; when it returns all requested + VMM-tagged allocations have been materialized on every rank. + Request admission reopens once every tag from the corresponding + sleep transition has been restored. Raises: ValueError: If the backend is not ``"pytorch"`` or @@ -934,13 +1013,28 @@ def wakeup(self, wakeup_tags: list[str]) -> None: tags = [ExecutorMemoryType(tag) for tag in wakeup_tags] logger.info(f"Wakeup: {tags}") - if self.llm_args.parallel_config.world_size > 1: - self._multi_rank_sleep_wakeup("wakeup", tags) - else: - with self.engine._sleep_wakeup_lock, self.engine.control_action(): - torch.cuda.synchronize() - materialize_with_tag(*tags) - torch.cuda.synchronize() + self.engine.begin_wakeup_transition(tags) + local_mutation_started = False + try: + if self.llm_args.parallel_config.world_size > 1: + self._multi_rank_sleep_wakeup("wakeup", tags) + else: + with self.engine._sleep_wakeup_lock, self.engine.control_action( + ): + torch.cuda.synchronize() + local_mutation_started = True + materialize_with_tag(*tags) + torch.cuda.synchronize() + except Exception: + if local_mutation_started: + self.engine.fail_sleep_wakeup_transition() + self.engine.abort_wakeup_transition() + raise + try: + self.engine.complete_wakeup_transition() + except Exception: + self.engine.fail_sleep_wakeup_transition() + raise def shutdown(self): if self.doing_shutdown: @@ -948,9 +1042,13 @@ def shutdown(self): else: self.doing_shutdown = True - if self.engine is not None and self.engine.can_enqueue_requests(): - self.engine.shutdown() - self.engine = None + if self.engine is not None: + can_shutdown = getattr(self.engine, "can_shutdown", None) + if can_shutdown is None: + can_shutdown = self.engine.can_enqueue_requests + if can_shutdown(): + self.engine.shutdown() + self.engine = None def get_disaggregated_params(self) -> dict: if self.engine is None or self.engine.kv_cache_transceiver is None: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5b39eea23efa..c2ed366ccf46 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -18,6 +18,8 @@ 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 - unittest/_torch/modeling/test_modeling_pixtral.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c605b8e36d6e..effeab272165 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -70,6 +70,7 @@ l0_cpu: - unittest/executor/test_multi_frontend_routing.py - unittest/executor/test_event_loop_error_broadcast.py - unittest/executor/test_rpc_worker_mixin.py + - unittest/executor/test_sleep_collective_rpc_guards.py - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_ray_stub.py diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index dceecb3761be..f3597fbe7a9a 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -55,6 +55,9 @@ l0_gb200_multi_gpus: - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_fabric_memory[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_bounce[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/unittest/_torch/executor/test_executor_request_queue.py b/tests/unittest/_torch/executor/test_executor_request_queue.py index 1478cb4ab990..d4b46e4502d8 100644 --- a/tests/unittest/_torch/executor/test_executor_request_queue.py +++ b/tests/unittest/_torch/executor/test_executor_request_queue.py @@ -16,7 +16,8 @@ import pytest from tensorrt_llm._torch.pyexecutor.executor_request_queue import ( - SHUTDOWN_REQUEST_ID, ExecutorRequestQueue, RequestQueueItem) + SHUTDOWN_REQUEST_ID, ExecutorRequestQueue, RequestAdmissionState, + RequestQueueItem) pytestmark = pytest.mark.cpu_only @@ -65,6 +66,7 @@ def test_executor_queue_init(executor_queue, mock_dist): assert executor_queue.next_request_id == 8 assert executor_queue.enable_iter_perf_stats assert executor_queue.active + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING assert isinstance(executor_queue.request_queue, queue.Queue) assert isinstance(executor_queue.enqueue_lock, type(threading.Lock())) @@ -155,6 +157,69 @@ def test_enqueue_request_after_shutdown(executor_queue): executor_queue.enqueue_request(Mock()) +def test_sleep_wakeup_admission_state_rejects_direct_enqueue(executor_queue): + """Admission stays closed from PARKING until the matching wakeup completes.""" + executor_queue.begin_sleep_transition(["model", "kv_cache"]) + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKING + assert not executor_queue.can_enqueue_request() + with pytest.raises(RuntimeError, match="admission is parking"): + executor_queue.enqueue_request(Mock()) + + executor_queue.complete_sleep_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + assert executor_queue.can_enqueue_control_request() + with pytest.raises(RuntimeError, match="admission is parked"): + executor_queue.enqueue_request(Mock()) + + executor_queue.begin_wakeup_transition(["kv_cache"]) + assert executor_queue.get_admission_state() is RequestAdmissionState.WAKING + with pytest.raises(RuntimeError, match="admission is waking"): + executor_queue.enqueue_request(Mock()) + + executor_queue.complete_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + with pytest.raises(RuntimeError, match="admission is parked"): + executor_queue.enqueue_request(Mock()) + + executor_queue.begin_wakeup_transition(["model"]) + executor_queue.complete_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + assert executor_queue.can_enqueue_request() + + +def test_wakeup_allows_extra_tags_and_reopens_after_parked_tags(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.complete_sleep_transition() + + executor_queue.begin_wakeup_transition(["model", "kv_cache"]) + executor_queue.complete_wakeup_transition() + + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + + +def test_recoverable_sleep_wakeup_failures_restore_prior_state(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.abort_sleep_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + + executor_queue.begin_sleep_transition(["model"]) + executor_queue.complete_sleep_transition() + executor_queue.begin_wakeup_transition(["model"]) + executor_queue.abort_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + + +def test_post_mutation_failure_permanently_closes_admission(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.fail_sleep_wakeup_transition() + executor_queue.abort_sleep_transition() + + assert executor_queue.get_admission_state() is RequestAdmissionState.FAILED + assert not executor_queue.can_enqueue_request() + with pytest.raises(RuntimeError, match="admission is failed"): + executor_queue.enqueue_request(Mock()) + + @pytest.mark.parametrize( "rank,active,expected", [ diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 3919ea233af4..92bb2544ed6a 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -17,12 +17,18 @@ import sys from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from tensorrt_llm._torch.modules.fused_moe import nccl_ep_utils from tensorrt_llm._torch.modules.fused_moe.communication import communication_factory +from tensorrt_llm._torch.modules.fused_moe.communication import nvlink_one_sided as one_sided_module +from tensorrt_llm._torch.modules.fused_moe.communication import nvlink_two_sided as two_sided_module +from tensorrt_llm._torch.modules.fused_moe.communication import ( + nvlink_two_sided_flashinfer as flashinfer_module, +) from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) @@ -39,6 +45,7 @@ def _make_model_config( moe_tp_size=1, moe_ep_size=2, moe_ep_rank=0, + has_cp_helix=Mock(return_value=False), ) return SimpleNamespace( mapping=mapping, @@ -57,6 +64,116 @@ def _strategy_unavailable(*args, **kwargs): raise RuntimeError("strategy unavailable") +@pytest.mark.parametrize("use_flashinfer", [False, True]) +def test_forced_two_sided_rejects_helix_before_workspace_allocation( + monkeypatch: pytest.MonkeyPatch, + use_flashinfer: bool, +) -> None: + model_config = _make_model_config() + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkTwoSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSidedFlashinfer, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) + monkeypatch.setattr(two_sided_module.MnnvlMemory, "initialize", native_initialize) + monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) + + with pytest.raises(ValueError, match="does not support Helix context parallelism"): + communication_factory.CommunicationFactory._create_forced_method( + "NVLINK_TWO_SIDED", + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + payload_in_workspace=False, + alltoall_result_do_sum=True, + use_flashinfer=use_flashinfer, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + flashinfer_symbols.assert_not_called() + + +def test_forced_one_sided_rejects_helix_before_workspace_allocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_config = _make_model_config() + model_config.mapping.world_size = model_config.mapping.moe_ep_size + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkOneSided, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + monkeypatch.setattr(one_sided_module.MnnvlMemory, "initialize", native_initialize) + + with pytest.raises(ValueError, match="does not support Helix context parallelism"): + communication_factory.CommunicationFactory._create_forced_method( + "NVLINK_ONE_SIDED", + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + payload_in_workspace=False, + alltoall_result_do_sum=True, + use_flashinfer=False, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + + +def test_auto_selection_rejects_unrepurposed_helix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_config = _make_model_config() + model_config.mapping.world_size = model_config.mapping.moe_ep_size + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkOneSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSidedFlashinfer, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) + monkeypatch.setattr(one_sided_module.MnnvlMemory, "initialize", native_initialize) + monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) + with pytest.raises(ValueError, match="repurpose_helix_cp_to_tp"): + communication_factory.CommunicationFactory.create_strategy( + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + flashinfer_symbols.assert_not_called() + + def _install_failing_nccl_module(monkeypatch: pytest.MonkeyPatch, error: BaseException): def fail_get_version(): raise error diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index 8b564563a92e..757d06bfe90a 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -206,6 +206,7 @@ def test_communication_factory_accepts_model_selected_method(monkeypatch): dp_size=16, moe_tp_size=1, moe_ep_size=16, + has_cp_helix=lambda: False, ) model_config = SimpleNamespace( mapping=mapping, diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 4c3bbb84ce79..3ec329ca75cf 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -50,9 +50,11 @@ import os import pickle import sys +import threading import traceback from dataclasses import dataclass from functools import lru_cache +from types import SimpleNamespace from typing import Dict, List, Optional, Set, Tuple from unittest.mock import MagicMock @@ -62,7 +64,8 @@ from mpi4py import MPI import tensorrt_llm as tllm -from tensorrt_llm._mnnvl_utils import MnnvlMemory +import tensorrt_llm._mnnvl_utils as mnnvl +from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) @@ -2559,6 +2562,367 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" +def _exercise_mnnvl_checkpoint_graph_replay( + config: CommTestConfig, rank: int, communication +) -> bool: + worker_inputs = _prepare_worker_inputs(rank, config) + _, local_slot_start, local_slot_end = _compute_ep_partition( + config.num_experts, config.ep_size, rank + ) + + def forward() -> torch.Tensor: + dispatch_outputs = _run_worker_dispatch(communication, worker_inputs, config) + received_hidden_states = _to_bf16( + dispatch_outputs.recv_hs, + dispatch_outputs.recv_sf, + worker_inputs.global_scale, + config.quant_mode, + ) + moe_output = simple_moe( + received_hidden_states, + dispatch_outputs.recv_slots, + dispatch_outputs.recv_scales, + local_slot_start, + local_slot_end, + ) + return communication.combine( + moe_output, + all_rank_max_num_tokens=max(config.all_num_tokens), + ) + + with torch.inference_mode(): + forward() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = forward() + graph.replay() + torch.cuda.synchronize() + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + stable_addresses = (communication.mnnvl_mem.ptr, communication.workspace.data_ptr()) + + def allocations_are_mapped() -> bool: + return communication.mnnvl_mem.mapped + + else: + stable_addresses = ( + communication.alltoall_workspace.data_ptr(), + communication.alltoall_prepare_workspace.data_ptr(), + ) + + def allocations_are_mapped() -> bool: + return all( + workspace is None or workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) + + def any_allocation_is_mapped() -> bool: + return any( + workspace is not None and workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + + def any_allocation_is_mapped() -> bool: + return communication.mnnvl_mem.mapped + + for cycle in range(3): + free_before_prepare = torch.cuda.mem_get_info()[0] + communication.checkpoint_prepare() + assert not any_allocation_is_mapped() + fresh_comm = MPI.COMM_WORLD.Split(0, rank) + communication.checkpoint_restore(fresh_comm) + assert allocations_are_mapped() + memory_drift_bytes = free_before_prepare - torch.cuda.mem_get_info()[0] + assert memory_drift_bytes <= 64 * 1024**2 + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + restored_addresses = ( + communication.mnnvl_mem.ptr, + communication.workspace.data_ptr(), + ) + else: + restored_addresses = ( + communication.alltoall_workspace.data_ptr(), + communication.alltoall_prepare_workspace.data_ptr(), + ) + assert restored_addresses == stable_addresses + + with torch.inference_mode(): + worker_inputs.hs.add_(cycle + 1) + expected = forward().clone() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected) + + return True + + +def _worker_mnnvl_checkpoint_graph_replay(config: CommTestConfig) -> bool: + """Exercise stable-VA checkpoint cycles through a captured MoE communication graph.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + communication = create_comm_object(config.comm_type, mapping, config) + try: + return _exercise_mnnvl_checkpoint_graph_replay(config, rank, communication) + finally: + if hasattr(communication, "destroy"): + communication.destroy() + + +def _mnnvl_workspace_is_mapped(communication) -> bool: + if isinstance(communication, NVLinkOneSided): + return communication.mnnvl_mem.mapped + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + assert any(workspace is not None for workspace in workspaces) + return all(workspace is None or workspace.mapped for workspace in workspaces) + + +def _worker_mnnvl_engine_checkpoint_coordination(config: CommTestConfig) -> bool: + """Join the production engine coordinator, listener, and native resource hooks.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + rank = tllm.mpi_rank() + world_size = config.ep_size + torch.cuda.set_device(rank) + assert MPI.Query_thread() >= MPI.THREAD_MULTIPLE + + mapping = Mapping( + rank=rank, + tp_size=world_size, + moe_ep_size=world_size, + world_size=world_size, + ) + communication = create_comm_object(config.comm_type, mapping, config) + sleep_wakeup_comm = MPI.COMM_WORLD.Dup() + control_comm = MPI.COMM_WORLD.Dup() + + class _Model: + def modules(self): + return [SimpleNamespace(comm=communication)] + + executor = object.__new__(PyExecutor) + executor._sleep_wakeup_comm = sleep_wakeup_comm + executor._sleep_wakeup_lock = threading.Lock() + executor._sleep_wakeup_listener_thread = None + executor.control_request_barrier = threading.Event() + executor.control_action_done = threading.Event() + executor._active_control_id = None + executor.device_id = rank + executor.dist = SimpleNamespace(rank=rank, world_size=world_size) + executor.model_engine = SimpleNamespace(model=_Model()) + executor.draft_model_engine = None + + class _ControlRequestQueue: + def enqueue_control_request(self, *, drain, control_id): + assert drain + control_comm.bcast(control_id, root=0) + executor._active_control_id = control_id + executor.control_action_done.clear() + executor.control_request_barrier.set() + + executor.executor_request_queue = _ControlRequestQueue() + retained_objects = getattr(sys, "_trtllm_mnnvl_coordination_test_objects", None) + if retained_objects is None: + retained_objects = [] + setattr(sys, "_trtllm_mnnvl_coordination_test_objects", retained_objects) + retained_objects.append((communication, sleep_wakeup_comm, control_comm, executor)) + worker = object.__new__(BaseWorker) + worker.doing_shutdown = True + worker._backend = "pytorch" + worker.rank = rank + worker.engine = executor + worker.llm_args = SimpleNamespace( + parallel_config=SimpleNamespace(world_size=world_size), + sleep_config=object(), + ) + + listener = None + listener_shutdown = False + try: + initial_error = None + try: + assert executor._has_mnnvl_checkpoint_resources([ExecutorMemoryType.MODEL_ENGINE_MAIN]) + assert _mnnvl_workspace_is_mapped(communication) + except Exception as error: + initial_error = f"rank {rank}: {error}" + initial_errors = MPI.COMM_WORLD.allgather(initial_error) + assert not any(initial_errors), initial_errors + + if rank != 0: + listener = threading.Thread( + target=executor._sleep_wakeup_listener_loop, + name=f"mnnvl-checkpoint-listener-{rank}", + daemon=True, + ) + executor._sleep_wakeup_listener_thread = listener + listener.start() + + MPI.COMM_WORLD.Barrier() + expected_mapped_states = (False, True) + if rank == 0: + for action, expected_mapped in zip(("sleep", "wakeup"), expected_mapped_states): + action_error = None + try: + worker._multi_rank_sleep_wakeup(action, [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + except Exception as error: + action_error = error + control_comm.bcast(action_error is None, root=0) + if action_error is not None: + raise action_error + + local_state_error = None + try: + assert _mnnvl_workspace_is_mapped(communication) is expected_mapped + except Exception as error: + local_state_error = f"rank {rank}: {error}" + state_errors = control_comm.allgather(local_state_error) + assert not any(state_errors), state_errors + executor._shutdown_sleep_wakeup_listeners() + listener_shutdown = True + else: + # This is the only emulated production piece: the regular request + # broadcaster delivers the control sentinel to the executor loop. + # The real listener and its dedicated communicator remain under test. + for expected_mapped in expected_mapped_states: + control_id = control_comm.bcast(None, root=0) + executor._active_control_id = control_id + executor.control_action_done.clear() + executor.control_request_barrier.set() + action_completed = executor.control_action_done.wait(timeout=60.0) + executor.control_action_done.clear() + executor._active_control_id = None + action_succeeded = control_comm.bcast(None, root=0) + if not action_succeeded: + break + + local_state_error = None + try: + assert action_completed, "listener did not release the control request" + assert _mnnvl_workspace_is_mapped(communication) is expected_mapped + except Exception as error: + local_state_error = f"rank {rank}: {error}" + state_errors = control_comm.allgather(local_state_error) + assert not any(state_errors), state_errors + assert listener is not None + listener.join(timeout=60.0) + assert not listener.is_alive() + return True + finally: + if rank == 0 and not listener_shutdown: + executor._shutdown_sleep_wakeup_listeners() + if listener is not None and listener.is_alive(): + listener.join(timeout=1.0) + # The process-global retention list keeps the duplicated communicators, + # executor shell, and native resource alive until MPI_Finalize. Freeing + # or destroying any of them while a failed peer listener still owns a + # request can hide the original failure with a hang/crash. + + +class _FailureInjectionMnnvlMemory(MnnvlMemory): + pass + + +def _worker_mnnvl_checkpoint_failure_injection(_unused=None) -> bool: + """Prove local CUDA failures do not strand peers in checkpoint collectives.""" + from types import SimpleNamespace + from unittest.mock import Mock, patch + + rank = tllm.mpi_rank() + comm = MPI.COMM_WORLD.Split(0, rank) + comm_size = comm.Get_size() + obj = _FailureInjectionMnnvlMemory.__new__(_FailureInjectionMnnvlMemory) + obj.ptr = 1032 + obj.mapping = SimpleNamespace(rank=rank) + + def install_record(state, handles): + record = mnnvl._MnnvlAllocationRecord( + comm=comm, + comm_size=comm_size, + comm_rank=rank, + comm_membership=tuple(range(comm_size)), + aligned_size=64, + mem_handles=handles, + start_address=1000, + rank_stride=256, + address_offset=32, + state=state, + ) + _FailureInjectionMnnvlMemory.allocated_map = {obj.ptr: record} + return record + + prepare_record = install_record( + mnnvl._MnnvlAllocationState.MAPPED, + list(range(1, comm_size + 1)), + ) + unmap = Mock() + if rank == 0: + unmap.side_effect = [ + None, + RuntimeError("injected partial unmap failure"), + *([None] * (comm_size - 2)), + ] + with ( + patch.object(mnnvl.torch.cuda, "synchronize"), + patch.object(mnnvl.cuda, "cuMemUnmap", unmap), + patch.object(mnnvl.cuda, "cuMemRelease", return_value=None), + patch.object(mnnvl, "_check_cu_result", side_effect=lambda result: result), + ): + prepare_error = None + try: + obj.checkpoint_prepare() + except RuntimeError as error: + prepare_error = str(error) + + prepare_errors = comm.allgather(prepare_error) + assert "injected partial unmap failure" in prepare_errors[0] + assert all(error is None for error in prepare_errors[1:]) + assert prepare_record.state is ( + mnnvl._MnnvlAllocationState.BROKEN if rank == 0 else mnnvl._MnnvlAllocationState.UNMAPPED + ) + + restore_record = install_record( + mnnvl._MnnvlAllocationState.UNMAPPED, + [None] * comm_size, + ) + create_and_map = Mock( + side_effect=(RuntimeError("injected remap failure") if rank == 0 else None), + return_value=list(range(11, 11 + comm_size)), + ) + with ( + patch.object(mnnvl.torch.cuda, "synchronize"), + patch.object(mnnvl.cuda, "cuMemUnmap", return_value=None), + patch.object(mnnvl.cuda, "cuMemRelease", return_value=None), + patch.object(mnnvl, "_check_cu_result", side_effect=lambda result: result), + patch.object( + _FailureInjectionMnnvlMemory, + "_create_and_map_handles", + create_and_map, + ), + ): + restore_error = None + try: + obj.checkpoint_restore(comm) + except RuntimeError as error: + restore_error = str(error) + + restore_errors = comm.allgather(restore_error) + assert all("injected remap failure" in error for error in restore_errors) + assert restore_record.state is mnnvl._MnnvlAllocationState.BROKEN + _FailureInjectionMnnvlMemory.allocated_map = {} + del obj.ptr + return True + + # ============================================================================ # Test Class # ============================================================================ @@ -2624,6 +2988,77 @@ def test_nccl_ep_cuda_graph_replay_uses_updated_routing(self, mpi_pool_executor) """Verify LL CUDA graph replay reads routing written after capture.""" _run_nccl_ep_cuda_graph_replay_test(mpi_pool_executor) + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + @pytest.mark.parametrize( + "comm_type", + [COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED], + ) + def test_mnnvl_checkpoint_preserves_moe_graph_addresses( + self, + mpi_pool_executor, + comm_type: str, + ) -> None: + """Verify repeated checkpoint cycles preserve MoE graph pointers and outputs.""" + config = CommTestConfig( + comm_type=comm_type, + ep_size=mpi_pool_executor.num_workers, + num_experts=32, + top_k=2, + hidden_size=128, + all_num_tokens=[8] * mpi_pool_executor.num_workers, + ) + skip_reason = _get_skip_reason(config) + if skip_reason: + pytest.skip(skip_reason) + results = mpi_pool_executor.map( + _worker_mnnvl_checkpoint_graph_replay, + [config] * config.ep_size, + ) + assert all(results) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + @pytest.mark.parametrize( + "comm_type", + [COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED], + ) + def test_mnnvl_engine_checkpoint_coordination( + self, + mpi_pool_executor, + comm_type: str, + ) -> None: + """Exercise BaseWorker and PyExecutor around native MNNVL checkpoint hooks.""" + config = CommTestConfig( + comm_type=comm_type, + ep_size=mpi_pool_executor.num_workers, + num_experts=32, + top_k=2, + hidden_size=128, + all_num_tokens=[8] * mpi_pool_executor.num_workers, + ) + skip_reason = _get_skip_reason(config) + if skip_reason: + pytest.skip(skip_reason) + results = mpi_pool_executor.map( + _worker_mnnvl_engine_checkpoint_coordination, + [config] * config.ep_size, + ) + assert all(results) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + def test_mnnvl_checkpoint_failure_is_collective_and_bounded( + self, + mpi_pool_executor, + ) -> None: + """A single-rank unmap or remap failure must return on every rank.""" + results = mpi_pool_executor.map( + _worker_mnnvl_checkpoint_failure_injection, + [None] * mpi_pool_executor.num_workers, + ) + assert all(results) + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( "mpi_pool_executor,local_num_tokens,top_k", 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..fbd210fa9b87 --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -0,0 +1,881 @@ +# 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_timeout_fails_closed(monkeypatch) -> None: + lifecycle, memory, _ = _make_lifecycle() + request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + memory.comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda source, tag: request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + with pytest.raises(TimeoutError, match="workspace idle readiness"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_fail_closed.assert_called_once_with() + 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_prepare_timeout_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda source, tag: request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + ) + main_workspace = Mock(mapped=True, comm=comm) + prepare_workspace = Mock(mapped=True) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", main_workspace) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + prepare_workspace, + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + with pytest.raises(TimeoutError, match="workspace idle readiness"): + owner.checkpoint_prepare() + + main_workspace.checkpoint_fail_closed.assert_called_once_with() + prepare_workspace.checkpoint_fail_closed.assert_called_once_with() + + +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.can_use_cft_counted_writes = False + 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 + wrapper._workspace_registered = True + 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_moe_alltoall_aborted_registration_does_not_unregister( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle = Mock() + lifecycle.register.side_effect = RuntimeError("registration failed") + monkeypatch.setattr(MoeAlltoAll, "_WORKSPACES", {}) + monkeypatch.setattr(MoeAlltoAll, "_init_constants", Mock()) + monkeypatch.setattr( + MoeAlltoAll, + "_METAINFO_INDEX", + { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 0, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 0, + }, + ) + monkeypatch.setattr(mnnvl.MnnvlMemory, "initialize", Mock()) + memory = Mock() + memory.as_torch_strided_tensor.return_value = torch.zeros(1, dtype=torch.uint8) + monkeypatch.setattr( + "tensorrt_llm._torch.distributed.moe_alltoall.MnnvlMemory", + Mock(return_value=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(moe_ep_size=2, moe_ep_rank=0) + + with pytest.raises(RuntimeError, match="registration failed"): + MoeAlltoAll( + mapping=mapping, + max_num_tokens=1, + top_k=1, + num_slots=2, + workspace_size_per_rank=1, + ) + + lifecycle.unregister.assert_not_called() + + +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.can_use_cft_counted_writes = 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 + wrapper._workspace_registered = True + + 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, + has_cp_helix=Mock(return_value=False), + ) + + 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 + lifecycle.unregister.assert_not_called() diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py new file mode 100644 index 000000000000..5a20dde00f6a --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -0,0 +1,759 @@ +# 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 +import torch + +import tensorrt_llm._mnnvl_utils as mnnvl +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, membership=(0, 1)): + self.rank = rank + self.size = size + self.membership = membership + self.barrier_count = 0 + + def Get_rank(self): + return self.rank + + def Get_size(self): + return self.size + + def barrier(self): + self.barrier_count += 1 + + def allgather(self, value): + if isinstance(value, int) and not isinstance(value, bool): + return list(self.membership) + return [value] * self.size + + +class _TestMnnvlMemory(mnnvl.MnnvlMemory): + pass + + +def test_checkpoint_allgather_timeout_is_bounded(monkeypatch): + orphaned_requests = [] + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_ORPHANED_REQUESTS", orphaned_requests) + receive_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + send_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda **kwargs: receive_request, + isend=lambda value, **kwargs: send_request, + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + with pytest.raises(TimeoutError, match="mapping readiness"): + mnnvl._checkpoint_allgather( + comm, + None, + operation="mapping readiness", + ) + receive_request.Cancel.assert_called_once_with() + receive_request.Free.assert_not_called() + send_request.Cancel.assert_called_once_with() + send_request.Free.assert_not_called() + assert orphaned_requests == [receive_request, send_request] + + +def test_checkpoint_request_cleanup_retains_request_when_cancel_fails(monkeypatch): + orphaned_requests = [] + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_ORPHANED_REQUESTS", orphaned_requests) + request = SimpleNamespace( + test=Mock(return_value=(False, None)), + Cancel=Mock(side_effect=RuntimeError("cancel failed")), + Free=Mock(), + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + + mnnvl._cancel_checkpoint_requests([request]) + + request.Cancel.assert_called_once_with() + request.test.assert_called_once_with() + request.Free.assert_not_called() + assert orphaned_requests == [request] + + +def test_checkpoint_allgather_uses_bounded_object_requests(): + class _ReceiveRequest: + def __init__(self, value): + self.value = value + + def test(self): + return True, self.value + + comm = SimpleNamespace( + Get_rank=lambda: 1, + Get_size=lambda: 3, + irecv=lambda source, tag: _ReceiveRequest(f"rank-{source}"), + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + allgather=Mock(side_effect=AssertionError("blocking fallback used")), + ) + + result = mnnvl._checkpoint_allgather(comm, "rank-1", operation="test") + + assert result == ["rank-0", "rank-1", "rank-2"] + comm.allgather.assert_not_called() + + +def test_checkpoint_allgather_post_failure_retires_partial_requests(): + requests = [ + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + ] + receive_requests = iter(requests[:2]) + send_requests = iter([requests[2], RuntimeError("post failed")]) + + def isend(value, dest, tag): + result = next(send_requests) + if isinstance(result, Exception): + raise result + return result + + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 3, + irecv=lambda source, tag: next(receive_requests), + isend=isend, + ) + + with pytest.raises(RuntimeError, match="post failed"): + mnnvl._checkpoint_allgather(comm, None, operation="test") + + for request in requests: + request.Cancel.assert_called_once_with() + request.test.assert_called_once_with() + request.Free.assert_not_called() + + +@pytest.fixture +def memory(monkeypatch): + comm = _FakeComm() + record = mnnvl._MnnvlAllocationRecord( + comm=comm, + comm_size=2, + comm_rank=0, + comm_membership=(0, 1), + aligned_size=64, + mem_handles=[11, 22], + start_address=1000, + rank_stride=256, + address_offset=32, + ) + 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 + _TestMnnvlMemory.current_rank_stride = record.rank_stride + _TestMnnvlMemory.current_mem_offset = record.address_offset + record.aligned_size + + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", Mock()) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuMemUnmap", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + yield obj, record + _TestMnnvlMemory.allocated_map = {} + _TestMnnvlMemory.address_refcnt = {} + if hasattr(obj, "ptr"): + del obj.ptr + + +def test_checkpoint_prepare_preserves_va_and_is_idempotent(memory): + obj, record = memory + + obj.checkpoint_prepare() + + assert not obj.mapped + assert record.start_address == 1000 + assert record.rank_stride == 256 + assert record.address_offset == 32 + assert record.mem_handles == [None, None] + assert record.comm.barrier_count == 0 + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list] == [1032, 1288] + + obj.checkpoint_prepare() + assert record.comm.barrier_count == 0 + + +def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + restored_comm = _FakeComm() + create_and_map = Mock(return_value=[33, 44]) + monkeypatch.setattr(_TestMnnvlMemory, "_create_and_map_handles", create_and_map) + + assert obj.checkpoint_restore(restored_comm) + + create_and_map.assert_called_once_with(restored_comm, 64, 1000, 256, 32) + assert obj.ptr == 1032 + 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_checkpoint_restore_membership_timeout_is_terminal(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + receive_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda source, tag: receive_request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + ) + with pytest.raises(TimeoutError, match="communicator membership"): + obj.checkpoint_restore(comm) + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.pending_comm is None + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_restore(comm) + + +def test_checkpoint_restore_timeout_cleans_locally_mapped_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + pending = object() + + class Request: + def __init__(self, result=pending): + self.result = result + + def test(self): + return (self.result is not pending, None if self.result is pending else self.result) + + def Cancel(self): + pass + + def Free(self): + pass + + class TimeoutComm(_FakeComm): + def __init__(self): + super().__init__() + self.receive_requests = iter([Request(1), Request()]) + + def irecv(self, source, tag): + return next(self.receive_requests) + + def isend(self, value, dest, tag): + return Request(None) + + with pytest.raises(TimeoutError, match="mapping readiness"): + obj.checkpoint_restore(TimeoutComm()) + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [None, None] + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list[-2:]] == [1032, 1288] + + +def test_handle_exchange_timeout_does_not_start_second_collective(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(side_effect=TimeoutError("handle exchange timed out")), + ) + + class CountingComm(_FakeComm): + def __init__(self): + super().__init__() + self.allgather_count = 0 + + def allgather(self, value): + self.allgather_count += 1 + return super().allgather(value) + + comm = CountingComm() + with pytest.raises(TimeoutError, match="handle exchange timed out"): + obj.checkpoint_restore(comm) + + assert comm.allgather_count == 1 + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + + +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() + + with pytest.raises(RuntimeError, match="rank/size 1/2 != 0/2"): + obj.checkpoint_restore(_FakeComm(rank=1)) + + 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 == 0 + first._checkpoint_restore_complete.assert_called_once_with() + second._checkpoint_restore_complete.assert_called_once_with() + + +def test_mnnvl_moe_restore_prepare_only_skips_main_workspace_initialization(monkeypatch): + main_workspace = Mock() + main_workspace.checkpoint_restore.return_value = False + prepare_workspace = Mock() + prepare_workspace.checkpoint_restore.return_value = True + workspace_tensor = Mock() + initialize_workspace = Mock() + synchronize = Mock() + comm = _FakeComm() + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", main_workspace) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", prepare_workspace) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace_tensor", workspace_tensor) + monkeypatch.setattr( + mnnvl.torch.ops.trtllm, + "moe_initialize_workspace", + initialize_workspace, + ) + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", synchronize) + + mnnvl.MnnvlMoe.checkpoint_restore(comm) + + initialize_workspace.assert_not_called() + synchronize.assert_called_once_with() + assert comm.barrier_count == 0 + main_workspace._checkpoint_restore_complete.assert_not_called() + prepare_workspace._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() + address_free = Mock(return_value=None) + monkeypatch.setattr(mnnvl.cuda, "CUdeviceptr", lambda value: value) + monkeypatch.setattr(mnnvl.cuda, "cuMemAddressFree", address_free) + + _TestMnnvlMemory.close_mnnvl_memory(obj.ptr) + + address_free.assert_called_once_with( + record.start_address, record.comm_size * record.rank_stride + ) + assert obj.ptr not in _TestMnnvlMemory.allocated_map + assert record.start_address not in _TestMnnvlMemory.address_refcnt + del obj.ptr + + +def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): + comm = Mock() + comm.isend = None + comm.irecv = None + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = lambda value: [ + value, + { + "error": None, + "handle": b"remote", + "is_fabric": True, + "pid": 1001, + }, + ] + allocation_prop = SimpleNamespace( + requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + location=object(), + ) + access_desc = SimpleNamespace(location=None, flags=None) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr( + mnnvl.cuda, + "cuMemExportToShareableHandle", + Mock(return_value=SimpleNamespace(data=b"local")), + ) + monkeypatch.setattr(mnnvl.cuda, "CUmemAccessDesc", Mock(return_value=access_desc)) + monkeypatch.setattr(mnnvl.cuda, "cuMemImportFromShareableHandle", Mock(return_value=22)) + map_memory = Mock(side_effect=[None, RuntimeError("map failed")]) + unmap_memory = Mock(return_value=None) + release_memory = Mock(return_value=None) + monkeypatch.setattr(mnnvl.cuda, "cuMemMap", map_memory) + monkeypatch.setattr(mnnvl.cuda, "cuMemSetAccess", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.cuda, "cuMemUnmap", unmap_memory) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", release_memory) + + with pytest.raises(RuntimeError, match="map failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + unmap_memory.assert_called_once_with(1032, 64) + assert [call.args[0] for call in release_memory.call_args_list] == [11, 22] + + +def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypatch): + comm = Mock() + comm.isend = None + comm.irecv = None + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = lambda value: [ + value, + { + "error": None, + "handle": b"remote", + "is_fabric": True, + "pid": 1001, + }, + ] + allocation_prop = SimpleNamespace( + requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + location=object(), + ) + release_memory = Mock(return_value=None) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr( + mnnvl.cuda, + "cuMemExportToShareableHandle", + Mock(side_effect=RuntimeError("export failed")), + ) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", release_memory) + + with pytest.raises(RuntimeError, match="export failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + release_memory.assert_called_once_with(11) + + +@pytest.mark.parametrize( + ("errno_value", "expected_hint"), + [(mnnvl.errno.EPERM, "--cap-add=SYS_PTRACE"), (mnnvl.errno.ENOSYS, "requires Linux 5.6+")], +) +def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_value, expected_hint): + comm = Mock() + comm.isend = None + comm.irecv = None + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = lambda value: ( + [ + value, + { + "error": None, + "handle": 101, + "is_fabric": False, + "pid": 1001, + }, + ] + if isinstance(value, dict) + else [value, value] + ) + allocation_prop = SimpleNamespace( + requestedHandleTypes=( + mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ), + location=object(), + ) + syscall = Mock(side_effect=[10, 11, -1]) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr(mnnvl.cuda, "cuMemExportToShareableHandle", Mock(return_value=100)) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.ctypes, "CDLL", Mock(return_value=SimpleNamespace(syscall=syscall))) + monkeypatch.setattr(mnnvl.ctypes, "get_errno", Mock(return_value=errno_value)) + close_fd = Mock() + monkeypatch.setattr(mnnvl.os, "close", close_fd) + + with pytest.raises(RuntimeError, match=expected_hint): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + assert comm.allgather.call_count == 2 + assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 100] + + +def test_create_and_map_handles_close_failure_does_not_mask_original_error(monkeypatch): + comm = Mock() + comm.isend = None + comm.irecv = None + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = lambda value: ( + [ + value, + { + "error": None, + "handle": 101, + "is_fabric": False, + "pid": 1001, + }, + ] + if isinstance(value, dict) + else [value, value] + ) + allocation_prop = SimpleNamespace( + requestedHandleTypes=( + mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ), + location=object(), + ) + access_desc = SimpleNamespace(location=None, flags=None) + syscall = Mock(side_effect=[10, 11, 20, 21]) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr(mnnvl.cuda, "cuMemExportToShareableHandle", Mock(return_value=100)) + monkeypatch.setattr(mnnvl.cuda, "CUmemAccessDesc", Mock(return_value=access_desc)) + monkeypatch.setattr(mnnvl.cuda, "cuMemMap", Mock(side_effect=RuntimeError("map failed"))) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.ctypes, "CDLL", Mock(return_value=SimpleNamespace(syscall=syscall))) + close_fd = Mock(side_effect=[OSError("close failed"), None, None, None, None]) + monkeypatch.setattr( + mnnvl, + "os", + SimpleNamespace(close=close_fd, getpid=mnnvl.os.getpid, strerror=mnnvl.os.strerror), + ) + + with pytest.raises(RuntimeError, match="map failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + assert comm.allgather.call_count == 2 + assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 20, 21, 100] + + +def _make_moe_alltoall_for_lifecycle(): + obj = MoeAlltoAll.__new__(MoeAlltoAll) + obj._destroyed = True + obj.mnnvl_mem = Mock(mapped=True) + return obj + + +def test_moe_alltoall_rejects_unmapped_workspace(): + obj = _make_moe_alltoall_for_lifecycle() + obj.mnnvl_mem.mapped = False + + with pytest.raises(RuntimeError, match="workspace handles are unmapped"): + obj._require_mapped() + + +def test_two_sided_combine_requires_new_prepare_before_next_dispatch(monkeypatch): + communication = NVLinkTwoSided.__new__(NVLinkTwoSided) + communication._dispatch_state = { + "alltoall_info": object(), + "original_token_count": 1, + } + communication.alltoall_workspace = object() + communication.ep_rank = 0 + communication.ep_size = 1 + communication.top_k = 1 + communication.use_low_precision_combine = False + communication.alltoall_result_do_sum = True + monkeypatch.setattr(mnnvl.MnnvlMoe, "require_mapped", Mock()) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "mnnvl_moe_alltoallv_combine", + Mock(return_value=torch.ones(1, 1)), + ) + + communication.combine(torch.ones(1, 1)) + + with pytest.raises(ValueError, match=r"requires prepare_dispatch\(\) to be called first"): + communication.dispatch(None, None, None, None, [1]) diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index 4576f88ede6c..c93849de08ab 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -50,6 +50,15 @@ def _make_worker(backend="pytorch", world_size=1, sleep_config=_SLEEP_CONFIG_DEF parallel_config=SimpleNamespace(world_size=world_size), sleep_config=sleep_config, ) + w.engine = SimpleNamespace( + begin_sleep_transition=MagicMock(), + complete_sleep_transition=MagicMock(), + abort_sleep_transition=MagicMock(), + begin_wakeup_transition=MagicMock(), + complete_wakeup_transition=MagicMock(), + abort_wakeup_transition=MagicMock(), + fail_sleep_wakeup_transition=MagicMock(), + ) return w @@ -111,6 +120,28 @@ def test_multirank_dispatches_to_helper(self, method): call_action, call_tags = mock_helper.call_args[0] assert call_action == method assert len(call_tags) == 1 + getattr(w.engine, f"begin_{method}_transition").assert_called_once_with(call_tags) + getattr(w.engine, f"complete_{method}_transition").assert_called_once_with() + getattr(w.engine, f"abort_{method}_transition").assert_not_called() + + def test_multirank_recoverable_failure_restores_admission(self, method): + """A helper failure before mutation restores the prior admission state.""" + from unittest.mock import patch + + w = _make_worker(world_size=2) + with ( + patch.object( + w, + "_multi_rank_sleep_wakeup", + side_effect=RuntimeError("prepare failed"), + ), + pytest.raises(RuntimeError, match="prepare failed"), + ): + getattr(w, method)(["kv_cache"]) + + getattr(w.engine, f"abort_{method}_transition").assert_called_once_with() + getattr(w.engine, f"complete_{method}_transition").assert_not_called() + w.engine.fail_sleep_wakeup_transition.assert_not_called() def test_backend_checked_before_sleep_config(self, method): """Backend check fires even when sleep_config is also absent.""" @@ -426,6 +457,115 @@ def test_multiple_peer_errors_aggregated(self): assert "rank 2 OOM" in msg +class TestMnnvlSleepWakeupCoordination: + """Native MNNVL hooks enter COMMIT collectively and remain tag-scoped.""" + + def test_pyexecutor_discovers_shared_native_resource_once(self): + from types import SimpleNamespace + from unittest.mock import Mock + + from tensorrt_llm._torch.modules.fused_moe.communication.base import ( + CheckpointableCommunication, + ) + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + class FutureCheckpointCommunication: + def __init__(self, resource_key): + self.resource_key = resource_key + self.checkpoint_prepare = Mock() + self.checkpoint_restore = Mock() + + def checkpoint_resource_key(self): + return self.resource_key + + shared_resource_key = object() + resources = [] + modules = [] + for _ in range(2): + resource = FutureCheckpointCommunication(shared_resource_key) + assert isinstance(resource, CheckpointableCommunication) + resources.append(resource) + modules.append(SimpleNamespace(comm=resource)) + + executor = object.__new__(PyExecutor) + executor.model_engine = SimpleNamespace(model=SimpleNamespace(modules=lambda: modules)) + executor.draft_model_engine = None + + tags = [ExecutorMemoryType.MODEL_ENGINE_MAIN] + assert executor._mnnvl_checkpoint_resources(tags) == [resources[0]] + executor._run_mnnvl_checkpoint_resources(_SleepWakeupAction.SLEEP, tags) + resources[0].checkpoint_prepare.assert_called_once_with() + resources[1].checkpoint_prepare.assert_not_called() + assert not executor._has_mnnvl_checkpoint_resources([ExecutorMemoryType.KV_CACHE]) + + def test_mnnvl_commit_reaches_peer_before_rank_zero_hook(self): + from unittest.mock import patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [ + {"status": "ok", "has_mnnvl_resources": True}, + {"status": "ok"}, + ], + world_size=2, + ) + events = [] + original_send = worker.engine._sleep_wakeup_comm.send + + def record_send(payload, *args, **kwargs): + events.append(("send", payload["action"])) + return original_send(payload, *args, **kwargs) + + worker.engine._sleep_wakeup_comm.send = record_send + # Model partitions can differ: rank 0 may own no MoE layer while a + # peer does. The PREPARE ACK must still select peer-first COMMIT. + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: False + worker.engine._run_mnnvl_checkpoint_resources = lambda action, tags: events.append( + ("local_mnnvl", action) + ) + + with ( + patch( + "tensorrt_llm._torch.virtual_memory.release_with_tag", + side_effect=lambda *tags: events.append(("local_vmm", "sleep")), + ), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + commit_index = events.index(("send", _SleepWakeupAction.COMMIT)) + hook_index = events.index(("local_mnnvl", _SleepWakeupAction.SLEEP)) + vmm_index = events.index(("local_vmm", "sleep")) + assert commit_index < hook_index < vmm_index + + def test_kv_only_sleep_does_not_run_mnnvl_hook(self): + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker([{"status": "ok"}, {"status": "ok"}], world_size=2) + run_mnnvl = Mock() + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: False + worker.engine._run_mnnvl_checkpoint_resources = run_mnnvl + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + run_mnnvl.assert_not_called() + + class TestMultiRankSendFailureRecovery: """Partial rank-0 broadcast failures must not leave peer ACKs undrained.""" @@ -464,8 +604,9 @@ def recv(self, source, tag): if source == 1: return {"status": "ok", "op_id": self.op_id} return { - "status": "error", - "error": "rank 0 aborted sleep/wakeup before local execution", + "status": "ok", + "error": None, + "reason": "rank 0 aborted sleep/wakeup before local execution", "op_id": self.op_id, } @@ -486,7 +627,6 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) - with ( patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), @@ -499,7 +639,6 @@ def _noop_control_action(**kwargs): msg = str(exc_info.value) assert "simulated mid-broadcast send failure" in msg - assert "rank 0 aborted sleep/wakeup" in msg assert release.call_count == 0 assert send_calls == [ (_SleepWakeupAction.PREPARE, 1, _SleepWakeupTag.ACTION), @@ -573,6 +712,8 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop monkeypatch.setattr(py_executor, "_SLEEP_WAKEUP_ACK_TIMEOUT_S", 0.0) monkeypatch.setattr(py_executor, "_SLEEP_WAKEUP_ACK_POLL_INTERVAL_S", 0.0) @@ -587,6 +728,72 @@ def _noop_control_action(**kwargs): with pytest.raises(RuntimeError, match="timed out waiting"): w._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + fail_stop.assert_called_once() + + def test_abort_send_failure_fail_stops_prepared_worker(self): + """Rank 0 must not resume if a prepared peer cannot be aborted.""" + from unittest.mock import patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "error", "error": "prepare failed"}], + world_size=2, + ) + original_send = worker.engine._sleep_wakeup_comm.send + + def fail_abort(payload, *args, **kwargs): + original_send(payload, *args, **kwargs) + if payload["action"] == _SleepWakeupAction.ABORT: + raise RuntimeError("abort send failed") + + worker.engine._sleep_wakeup_comm.send = fail_abort + fail_stop = MagicMock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="abort send failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + release.assert_not_called() + fail_stop.assert_called_once() + + def test_abort_error_ack_fail_stops_prepared_worker(self): + """A genuine ABORT handler error leaves peer recovery inconclusive.""" + from unittest.mock import patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [ + {"status": "error", "error": "prepare failed"}, + {"status": "error", "error": "abort handler failed"}, + ], + world_size=2, + ) + fail_stop = MagicMock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="abort handler failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + release.assert_not_called() + fail_stop.assert_called_once() + def test_commit_send_failure_aborts_uncommitted_rank(self): """A prepared rank that misses COMMIT must receive ABORT to unblock.""" import threading @@ -640,6 +847,8 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, @@ -652,6 +861,7 @@ def _noop_control_action(**kwargs): w._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) release.assert_called_once() + fail_stop.assert_called_once() assert send_calls == [ (_SleepWakeupAction.PREPARE, 1, _SleepWakeupTag.ACTION), (_SleepWakeupAction.PREPARE, 2, _SleepWakeupTag.ACTION), @@ -667,6 +877,228 @@ def _noop_control_action(**kwargs): (1, _SleepWakeupTag.ACK), ] + def test_mnnvl_partial_commit_still_runs_rank_zero_bounded_phase(self): + """Once a peer sees MNNVL COMMIT, rank 0 must enter the local phase.""" + import threading + from contextlib import contextmanager + from unittest.mock import Mock, patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + class FakeComm: + def __init__(self): + self.op_id = None + self.ack_phases = {1: [], 2: []} + + def send(self, payload, dest, tag): + self.op_id = payload.get("op_id", self.op_id) + if payload["action"] == _SleepWakeupAction.COMMIT and dest == 2: + raise RuntimeError("injected MNNVL commit send failure") + self.ack_phases[dest].append(payload["action"]) + + def iprobe(self, source, tag): + return True + + def recv(self, source, tag): + return { + "status": "ok", + "op_id": self.op_id, + "phase": self.ack_phases[source].pop(0), + "has_mnnvl_resources": True, + } + + @contextmanager + def control_action(**kwargs): + yield None + + worker = object.__new__(BaseWorker) + worker._backend = "pytorch" + worker.rank = 0 + worker.llm_args = SimpleNamespace( + backend="pytorch", + parallel_config=SimpleNamespace(world_size=3), + sleep_config=object(), + ) + run_mnnvl = Mock() + worker.engine = SimpleNamespace( + _sleep_wakeup_lock=threading.Lock(), + _sleep_wakeup_comm=FakeComm(), + control_action=control_action, + _has_mnnvl_checkpoint_resources=lambda tags: True, + _run_mnnvl_checkpoint_resources=run_mnnvl, + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="commit send failure"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + run_mnnvl.assert_called_once_with( + _SleepWakeupAction.SLEEP, + [ExecutorMemoryType.MODEL_ENGINE_MAIN], + ) + release.assert_called_once_with(ExecutorMemoryType.MODEL_ENGINE_MAIN) + fail_stop.assert_called_once() + + def test_mnnvl_all_commit_send_errors_are_treated_as_uncertain_delivery(self): + """A send error cannot prove that no peer received MNNVL COMMIT.""" + import threading + from contextlib import contextmanager + from unittest.mock import Mock, patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + class FakeComm: + def __init__(self): + self.op_id = None + self.ack_phases = {1: []} + + def send(self, payload, dest, tag): + self.op_id = payload.get("op_id", self.op_id) + if payload["action"] == _SleepWakeupAction.COMMIT: + raise RuntimeError("injected MNNVL commit send failure") + self.ack_phases[dest].append(payload["action"]) + + def iprobe(self, source, tag): + return True + + def recv(self, source, tag): + return { + "status": "ok", + "op_id": self.op_id, + "phase": self.ack_phases[source].pop(0), + "has_mnnvl_resources": True, + } + + @contextmanager + def control_action(**kwargs): + yield None + + worker = object.__new__(BaseWorker) + worker._backend = "pytorch" + worker.rank = 0 + worker.llm_args = SimpleNamespace( + backend="pytorch", + parallel_config=SimpleNamespace(world_size=2), + sleep_config=object(), + ) + run_mnnvl = Mock() + worker.engine = SimpleNamespace( + _sleep_wakeup_lock=threading.Lock(), + _sleep_wakeup_comm=FakeComm(), + control_action=control_action, + _has_mnnvl_checkpoint_resources=lambda tags: True, + _run_mnnvl_checkpoint_resources=run_mnnvl, + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="commit send failure"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + run_mnnvl.assert_called_once_with( + _SleepWakeupAction.SLEEP, + [ExecutorMemoryType.MODEL_ENGINE_MAIN], + ) + release.assert_called_once_with(ExecutorMemoryType.MODEL_ENGINE_MAIN) + fail_stop.assert_called_once() + + def test_postcommit_non_runtime_error_fail_stops_worker(self): + """Every ordinary exception after COMMIT must take the fail-stop path.""" + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "ok"}, {"status": "ok"}], + world_size=2, + ) + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: True + worker.engine._run_mnnvl_checkpoint_resources = Mock( + side_effect=ValueError("invalid checkpoint state") + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="invalid checkpoint state"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + fail_stop.assert_called_once() + + def test_divergent_commit_fail_stops_worker(self): + """A potentially divergent operation poisons the worker before returning.""" + from unittest.mock import patch + + from tensorrt_llm.executor.base_worker import BaseWorker + + worker = object.__new__(BaseWorker) + worker._fatal_error = None + worker.engine = SimpleNamespace( + _fatal_error=None, + is_shutdown=False, + fail_sleep_wakeup_transition=MagicMock(), + ) + error = RuntimeError("divergent sleep state") + + with patch("tensorrt_llm._torch.pyexecutor.hang_detector.propagate_hard_kill") as hard_kill: + worker._fail_stop_divergent_sleep_wakeup(error) + + assert worker._fatal_error is error + assert worker.engine._fatal_error is error + assert worker.engine.is_shutdown + worker.engine.fail_sleep_wakeup_transition.assert_called_once_with() + hard_kill.assert_called_once_with() + + def test_precommit_local_failure_remains_recoverable(self): + """A failure before local mutation aborts peers without fail-stopping.""" + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "ok"}, {"status": "ok"}], + world_size=2, + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize", side_effect=RuntimeError("sync failed")), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="sync failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + fail_stop.assert_not_called() + class TestListenerUncaughtExceptionSendsErrorAck: """An exception that bypasses the narrow except clause must still produce an error ACK. @@ -683,6 +1115,7 @@ def test_uncaught_exception_sends_error_ack(self): because error_msg is still None when the exception bypasses the except clause — leaving rank-0 with an inconsistent view of the operation. """ + import threading from types import SimpleNamespace from unittest.mock import patch @@ -702,6 +1135,10 @@ def send(self, payload, dest, tag): executor._sleep_wakeup_comm = FakeComm() executor.device_id = 0 executor.dist = SimpleNamespace(rank=1) + executor.control_request_barrier = threading.Event() + executor.control_request_barrier.set() + executor.control_action_done = threading.Event() + executor._active_control_id = None with ( patch("torch.cuda.set_device"), @@ -714,23 +1151,8 @@ def send(self, payload, dest, tag): ), patch("torch.cuda.synchronize"), ): - # The loop receives one message then raises MemoryError which is - # not in the narrow except. We patch recv to raise StopIteration - # on the second call so the loop terminates cleanly in the test. - call_count = [0] - - def _recv_once(self_inner, source, tag): - call_count[0] += 1 - if call_count[0] > 1: - raise StopIteration - return {"action": "sleep", "tags": ["kv_cache"]} - - executor._sleep_wakeup_comm.recv = lambda source, tag: _recv_once(None, source, tag) - - try: + with pytest.raises(MemoryError, match="simulated OOM outside except list"): executor._sleep_wakeup_listener_loop() - except StopIteration: - pass assert sent_acks, "finally block must send an ACK even for uncaught exceptions" assert sent_acks[0]["status"] == "error", ( @@ -739,6 +1161,8 @@ def _recv_once(self_inner, source, tag): ) assert sent_acks[0]["error"] is not None assert "MemoryError" in sent_acks[0]["error"] + assert not executor.control_request_barrier.is_set() + assert executor.control_action_done.is_set() class TestListenerAbortAndShutdown: @@ -803,10 +1227,11 @@ def send(self, payload, dest, tag): "error": None, "op_id": op_id, "phase": _SleepWakeupAction.PREPARE, + "has_mnnvl_resources": False, } ] - def test_abort_unblocks_control_request_and_sends_error_ack(self): + def test_abort_unblocks_control_request_and_sends_success_ack(self): """Abort messages release the non-rank executor control barrier.""" import threading from unittest.mock import patch @@ -858,8 +1283,9 @@ def send(self, payload, dest, tag): assert executor.control_action_done.is_set() assert not executor.control_request_barrier.is_set() assert sent_acks - assert sent_acks[0]["status"] == "error" - assert "rank 0 send failed" in sent_acks[0]["error"] + assert sent_acks[0]["status"] == "ok" + assert sent_acks[0]["error"] is None + assert "rank 0 send failed" in sent_acks[0]["reason"] def test_abort_before_control_barrier_unblocks_later_control_request(self): """An early ABORT is recorded and later consumed by matching control.""" @@ -922,7 +1348,9 @@ def send(self, payload, dest, tag): assert executor.control_requests == [] assert not executor.control_request_barrier.is_set() assert sent_acks - assert sent_acks[0]["status"] == "error" + assert sent_acks[0]["status"] == "ok" + assert sent_acks[0]["error"] is None + assert "rank 0 send failed" in sent_acks[0]["reason"] def test_shutdown_sends_ack_before_listener_exits(self): """Shutdown messages are acknowledged so rank-0 can drain them.""" @@ -1015,6 +1443,8 @@ def test_local_failure_drains_all_peer_acks(self): # Two peers prepare ok, then abort ok after rank-0 local failure. responses = [{"status": "ok"}, {"status": "ok"}] w, recv_calls = _make_proto_worker(responses, world_size=3) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch( @@ -1036,6 +1466,7 @@ def test_local_failure_drains_all_peer_acks(self): f"{len(recv_calls)}; stale ACKs would corrupt the next " "sleep/wakeup call." ) + fail_stop.assert_called_once() def test_local_failure_plus_peer_error_both_reported(self): """When rank-0 fails locally and a peer also errors, both messages must appear. @@ -1053,6 +1484,8 @@ def test_local_failure_plus_peer_error_both_reported(self): {"status": "error", "error": "rank 2 also failed"}, ] w, _ = _make_proto_worker(responses) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch( @@ -1070,6 +1503,7 @@ def test_local_failure_plus_peer_error_both_reported(self): msg = str(exc_info.value) assert "rank 0 VMM fault" in msg assert "rank 2 also failed" in msg + fail_stop.assert_called_once() class TestSingleRankLockAcquired: @@ -1113,6 +1547,13 @@ def _noop_control_action(): w.engine = SimpleNamespace( _sleep_wakeup_lock=SpyLock(), control_action=_noop_control_action, + begin_sleep_transition=MagicMock(), + complete_sleep_transition=MagicMock(), + abort_sleep_transition=MagicMock(), + begin_wakeup_transition=MagicMock(), + complete_wakeup_transition=MagicMock(), + abort_wakeup_transition=MagicMock(), + fail_sleep_wakeup_transition=MagicMock(), ) with ( @@ -1126,6 +1567,41 @@ def _noop_control_action(): assert lock_entered, f"{method}() with world_size=1 did not acquire _sleep_wakeup_lock" + @pytest.mark.parametrize( + ("method", "mutation"), + [("sleep", "release_with_tag"), ("wakeup", "materialize_with_tag")], + ) + def test_post_mutation_failure_permanently_closes_admission(self, method, mutation): + """A single-rank error after VMM mutation starts is fail-closed.""" + import threading + from contextlib import contextmanager + from unittest.mock import patch + + w = _make_worker() + + @contextmanager + def _noop_control_action(): + yield None + + w.engine._sleep_wakeup_lock = threading.Lock() + w.engine.control_action = _noop_control_action + + with ( + patch(f"tensorrt_llm._torch.virtual_memory.{mutation}"), + patch( + "torch.cuda.synchronize", + side_effect=[None, RuntimeError("post-mutation sync failed")], + ), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + pytest.raises(RuntimeError, match="post-mutation sync failed"), + ): + getattr(w, method)(["kv_cache"]) + + w.engine.fail_sleep_wakeup_transition.assert_called_once_with() + getattr(w.engine, f"abort_{method}_transition").assert_called_once_with() + getattr(w.engine, f"complete_{method}_transition").assert_not_called() + # --------------------------------------------------------------------------- # GenerationExecutorProxy / GenerationExecutorRpcProxy collective_rpc()