From 8a49b856694def3ecb7cd17d99d4dc10fe390cc7 Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Thu, 17 Sep 2026 15:16:24 +0200 Subject: [PATCH 1/2] distributed: add experimental endpoint-transfer review prototype Snapshot of the current endpoint-transfer prototype for review in the author's fork. Native GPU performance parity and full end-to-end framework validation remain open; this is not an upstream-ready change. Assisted-by: OpenAI Codex Signed-off-by: Samuel Nordmann --- build_variables.bzl | 1 + test/distributed/test_transfer.py | 9211 +++++++++++++++++ test/dynamo/test_trace_rules.py | 1 + torch/_C/_distributed_transfer.pyi | 7 + torch/csrc/Module.cpp | 2 + .../distributed/transfer/python_bindings.cpp | 432 + .../distributed/transfer/python_bindings.h | 9 + torch/distributed/_transfer/README.md | 814 ++ torch/distributed/_transfer/__init__.py | 123 + torch/distributed/_transfer/_api.py | 8283 +++++++++++++++ torch/distributed/_transfer/_backend.py | 1034 ++ torch/distributed/_transfer/_errors.py | 62 + torch/distributed/_transfer/_reference.py | 724 ++ torch/distributed/_transfer/_types.py | 867 ++ 14 files changed, 21570 insertions(+) create mode 100644 test/distributed/test_transfer.py create mode 100644 torch/_C/_distributed_transfer.pyi create mode 100644 torch/csrc/distributed/transfer/python_bindings.cpp create mode 100644 torch/csrc/distributed/transfer/python_bindings.h create mode 100644 torch/distributed/_transfer/README.md create mode 100644 torch/distributed/_transfer/__init__.py create mode 100644 torch/distributed/_transfer/_api.py create mode 100644 torch/distributed/_transfer/_backend.py create mode 100644 torch/distributed/_transfer/_errors.py create mode 100644 torch/distributed/_transfer/_reference.py create mode 100644 torch/distributed/_transfer/_types.py diff --git a/build_variables.bzl b/build_variables.bzl index 9591da3dbed6d..10b166d65fd20 100644 --- a/build_variables.bzl +++ b/build_variables.bzl @@ -975,6 +975,7 @@ libtorch_python_core_sources = [ "torch/csrc/autograd/python_variable_indexing.cpp", "torch/csrc/PyObjectConversionPythonImpl.cpp", "torch/csrc/distributed/python_placement.cpp", + "torch/csrc/distributed/transfer/python_bindings.cpp", "torch/csrc/dynamo/python_compiled_autograd.cpp", "torch/csrc/dynamo/cache_entry.cpp", "torch/csrc/dynamo/cpp_shim.cpp", diff --git a/test/distributed/test_transfer.py b/test/distributed/test_transfer.py new file mode 100644 index 0000000000000..b1aaa457eb93f --- /dev/null +++ b/test/distributed/test_transfer.py @@ -0,0 +1,9211 @@ +import ast +import gc +from array import array +from pathlib import Path +import sys +import threading +import weakref +from unittest import mock + +import torch +from torch.distributed._transfer import ( + BACKEND_FACTORY_API_VERSION, + BackendFailureError, + BackendIndexSelection, + BackendNotificationBatch, + BackendRequestSlot, + BackendRegion, + BackendWork, + BoundTransfer, + BusyError, + Capabilities, + ClosedError, + CudaOrderingMode, + Endpoint, + EndpointIdentity, + EndpointMetadata, + ImportedRawMemory, + ImportedRegistration, + IndexSelection, + InvalidMetadataError, + InvalidRegionError, + NotificationCompletion, + NotificationBatch, + RawSpan, + RequestLocalNotificationSlot, + RegisteredMemory, + RegistrationMetadata, + TransferCancelledError, + TransferOp, + TransferSlot, + TransferTimeoutError, + ThreadMode, + UnsupportedError, + Work, + WorkState, + register_backend, +) +from torch.distributed._transfer import _api, _backend, _reference, _types +from torch.distributed._transfer._reference import _ReferenceBackend +from torch.testing._internal.common_utils import run_tests, TestCase + + +def _emulated_compiled_index_analyzer( + local_indices: object, + remote_indices: object, + destination_is_remote: bool, +) -> tuple[bytes, int, int, int, int, int, int, bool]: + local_bytes = bytes(memoryview(local_indices).cast("B")) + remote_bytes = bytes(memoryview(remote_indices).cast("B")) + packed = local_bytes + remote_bytes + values = memoryview(packed).cast("i") + local_count = len(local_bytes) // 4 + remote_count = len(remote_bytes) // 4 + local = values[:local_count] + remote = values[local_count:] + destination = remote if destination_is_remote else local + return ( + packed, + local_count, + remote_count, + min(local), + max(local), + min(remote), + max(remote), + len(set(destination)) != len(destination), + ) + + +class _PendingWork(BackendWork): + def __init__(self) -> None: + self._state = WorkState.PENDING + + @property + def state(self) -> WorkState: + return self._state + + @property + def error(self) -> BaseException | None: + return None + + def cancel(self) -> bool: + self._state = WorkState.CANCELLED + return True + + def close(self) -> None: + pass + + +class _BadErrorWork(_PendingWork): + @property + def state(self) -> WorkState: + return WorkState.FAILED + + @property + def error(self): + return "not an exception" + + +class _RegressingWork(_PendingWork): + def __init__(self) -> None: + self._queries = 0 + self._cancelled = False + + @property + def state(self) -> WorkState: + if self._cancelled: + return WorkState.CANCELLED + self._queries += 1 + return WorkState.PENDING if self._queries == 2 else WorkState.RUNNING + + def cancel(self) -> bool: + self._cancelled = True + return True + + +class _TransientStateErrorWork(_PendingWork): + def __init__(self) -> None: + self._queries = 0 + self._cancelled = False + + @property + def state(self) -> WorkState: + self._queries += 1 + if self._queries == 1: + raise RuntimeError("transient query failure") + if self._cancelled: + return WorkState.CANCELLED + return WorkState.RUNNING + + def cancel(self) -> bool: + self._cancelled = True + return True + + +class _CountingPendingWork(_PendingWork): + def __init__(self) -> None: + super().__init__() + self.state_queries = 0 + + @property + def state(self) -> WorkState: + self.state_queries += 1 + return self._state + + +class _InterruptingFailureWork(_PendingWork): + def __init__(self) -> None: + super().__init__() + self.failure = RuntimeError("provider failure") + self.error_queries = 0 + + @property + def state(self) -> WorkState: + return WorkState.FAILED + + @property + def error(self) -> BaseException: + self.error_queries += 1 + if self.error_queries == 1: + raise KeyboardInterrupt + return self.failure + + +class _PendingReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + kwargs["adopt_work"](_PendingWork()) + + +class _CountingPendingReferenceBackend(_ReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.last_work: _CountingPendingWork | None = None + + def submit(self, *args, **kwargs) -> None: + self.last_work = _CountingPendingWork() + kwargs["adopt_work"](self.last_work) + + +class _InterruptingFailureReferenceBackend(_ReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.last_work: _InterruptingFailureWork | None = None + + def submit(self, *args, **kwargs) -> None: + self.last_work = _InterruptingFailureWork() + kwargs["adopt_work"](self.last_work) + + +class _FusedSelectionWork(_PendingWork): + def __init__(self, state: WorkState) -> None: + super().__init__() + self._state = state + self.state_queries = 0 + self.bounded_calls: list[tuple[int, int | None, int]] = [] + self.next_bounded_state = WorkState.COMPLETED + self._error: BaseException | None = None + + @property + def state(self) -> WorkState: + self.state_queries += 1 + return self._state + + @property + def error(self) -> BaseException | None: + return self._error + + def poll_bounded( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState: + self.bounded_calls.append((max_polls, timeout_ns, timeout_check_interval)) + self._state = self.next_bounded_state + return self._state + + +class _FusedSelectionReferenceBackend(_ReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + fused_prevalidated_submit_poll=True, + fused_work_poll=True, + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.initial_state = WorkState.RUNNING + self.fused_calls: list[tuple[int, int | None, bytes | None]] = [] + self.last_work: _FusedSelectionWork | None = None + + def submit_prevalidated_and_poll( + self, + op, + plan_handle, + *, + selection, + notification, + max_polls, + timeout_ns, + adopt_work, + ) -> WorkState: + selection._validate_binding(op, plan_handle) + work = _FusedSelectionWork(self.initial_state) + self.last_work = work + adopt_work(work) + self.fused_calls.append((max_polls, timeout_ns, notification)) + return self.initial_state + + +class _BackgroundOnlyReferenceBackend(_ReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + background_progress=True, + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + def progress(self) -> None: + raise AssertionError("background-only backend must not be progressed") + + +class _AttachedOnlyReferenceBackend(_ReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + ) + + +class _BadErrorReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + kwargs["adopt_work"](_BadErrorWork()) + + +class _InterruptReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + raise KeyboardInterrupt + + +class _AdoptThenInterruptReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + kwargs["adopt_work"](_PendingWork()) + raise KeyboardInterrupt + + +class _RegressingReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + kwargs["adopt_work"](_RegressingWork()) + + +class _TransientStateErrorReferenceBackend(_ReferenceBackend): + def submit(self, *args, **kwargs) -> None: + kwargs["adopt_work"](_TransientStateErrorWork()) + + +class _InterruptAfterCloseReferenceBackend(_ReferenceBackend): + def close(self) -> None: + was_closed = self._closed + super().close() + if not was_closed: + raise KeyboardInterrupt + + +class _TestRequestSlot(BackendRequestSlot): + def __init__(self) -> None: + self._state: WorkState | None = None + self._error: BaseException | None = None + self.next_state = WorkState.COMPLETED + self.next_error: BaseException | None = None + self.start_count = 0 + self.recycle_count = 0 + self.close_count = 0 + self.closed = False + self.fail_close_once = False + self.raise_error_once: BaseException | None = None + self.start_entered: threading.Event | None = None + self.start_release: threading.Event | None = None + self.completion = object() + + @property + def state(self) -> WorkState | None: + if self.closed: + raise ClosedError("test request slot is closed") + return self._state + + @property + def error(self) -> BaseException | None: + if self.closed: + raise ClosedError("test request slot is closed") + if self.raise_error_once is not None: + error = self.raise_error_once + self.raise_error_once = None + raise error + return self._error + + @property + def completion_event(self) -> object: + return self.completion + + def start(self) -> None: + if self.closed: + raise ClosedError("test request slot is closed") + self.start_count += 1 + if self.start_entered is not None: + self.start_entered.set() + if self.start_release is not None and not self.start_release.wait(2): + raise RuntimeError("test request slot start timed out") + self._state = self.next_state + self._error = self.next_error + + def cancel(self) -> bool: + if self.closed: + raise ClosedError("test request slot is closed") + if self._state in { + WorkState.COMPLETED, + WorkState.FAILED, + WorkState.CANCELLED, + }: + return self._state is WorkState.CANCELLED + self._state = WorkState.CANCELLED + self._error = None + return True + + def recycle(self) -> None: + if self.closed: + raise ClosedError("test request slot is closed") + self.recycle_count += 1 + self._state = None + self._error = None + + def close(self) -> None: + self.close_count += 1 + if self.fail_close_once: + self.fail_close_once = False + raise RuntimeError("injected request-slot cleanup failure") + self.closed = True + + +class _UnacknowledgedStartRequestSlot(_TestRequestSlot): + @property + def state(self) -> WorkState | None: + return None + + def start(self) -> None: + raise KeyboardInterrupt + + +class _NormallyUnacknowledgedRequestSlot(_UnacknowledgedStartRequestSlot): + def start(self) -> None: + self.start_count += 1 + + +class _BlockingObservationRequestSlot(_TestRequestSlot): + def __init__(self) -> None: + super().__init__() + self.block_next_observation = False + self.observation_entered = threading.Event() + self.observation_release = threading.Event() + + @property + def state(self) -> WorkState | None: + if self.block_next_observation: + self.block_next_observation = False + self.observation_entered.set() + if not self.observation_release.wait(2): + raise RuntimeError("test request-slot observation timed out") + return super().state + + +_RETURN_CURRENT_SLOT_STATE = object() + + +class _FusedTestRequestSlot(_TestRequestSlot): + def __init__(self) -> None: + super().__init__() + self.fused_calls: list[tuple[int, int | None]] = [] + self.fused_result: object = _RETURN_CURRENT_SLOT_STATE + self.fused_error: BaseException | None = None + self.poll_states: list[WorkState | None] = [] + self.poll_count = 0 + self.bounded_calls: list[tuple[int, int | None, int]] = [] + self.bounded_results: list[object] = [] + self.bounded_error: BaseException | None = None + self.bounded_entered: threading.Event | None = None + self.bounded_release: threading.Event | None = None + + def poll_state(self) -> WorkState | None: + self.poll_count += 1 + if self.poll_states: + state = self.poll_states.pop(0) + if isinstance(state, WorkState): + self._state = state + return state + return self.state + + def start_and_poll( + self, + *, + max_polls: int, + timeout_ns: int | None, + ): + self.fused_calls.append((max_polls, timeout_ns)) + self.start() + if self.fused_error is not None: + raise self.fused_error + if self.fused_result is _RETURN_CURRENT_SLOT_STATE: + return self._state + return self.fused_result + + def poll_bounded( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ): + self.bounded_calls.append((max_polls, timeout_ns, timeout_check_interval)) + if self.bounded_entered is not None: + self.bounded_entered.set() + if self.bounded_release is not None and not self.bounded_release.wait(2): + raise RuntimeError("test bounded request-slot poll timed out") + if self.bounded_error is not None: + raise self.bounded_error + if self.bounded_results: + state = self.bounded_results.pop(0) + if isinstance(state, WorkState): + self._state = state + return state + return super().poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + + +class _NotificationFusedTestRequestSlot(_FusedTestRequestSlot): + def __init__(self) -> None: + super().__init__() + self.notification_starts: list[bytes | None] = [] + self.notification_fused_calls: list[tuple[bytes | None, int, int | None]] = [] + + def start_with_notification(self, notification: bytes | None) -> None: + self.notification_starts.append(notification) + self.start() + + def start_and_poll_with_notification( + self, + notification: bytes | None, + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + self.notification_fused_calls.append((notification, max_polls, timeout_ns)) + self.start() + if self.fused_error is not None: + raise self.fused_error + if self.fused_result is _RETURN_CURRENT_SLOT_STATE: + return self._state + return self.fused_result + + +class _TestExecutionDispatch: + _OPERATION_NAMES = frozenset( + { + "start", + "start_and_poll", + "poll_bounded", + "poll_state", + "cancel", + "recycle", + } + ) + + def __init__(self, slot: _FusedTestRequestSlot) -> None: + self.slot = slot + self.attribute_reads = {name: 0 for name in self._OPERATION_NAMES} + + def __getattribute__(self, name): + if name in object.__getattribute__(self, "_OPERATION_NAMES"): + reads = object.__getattribute__(self, "attribute_reads") + reads[name] += 1 + return object.__getattribute__(self, name) + + def start(self) -> None: + self.slot.start() + + def start_and_poll(self, *, max_polls: int, timeout_ns: int | None): + return self.slot.start_and_poll( + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + + def poll_bounded( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ): + return self.slot.poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + + def poll_state(self) -> WorkState | None: + return self.slot.poll_state() + + def cancel(self) -> bool: + return self.slot.cancel() + + def recycle(self) -> None: + self.slot.recycle() + + @property + def error(self): + raise AssertionError("execution dispatch must not own diagnostics") + + @property + def completion_event(self): + raise AssertionError("execution dispatch must not own completion events") + + def close(self) -> None: + raise AssertionError("execution dispatch must not own close") + + +class _CapabilityGatedExecutionDispatch(_TestExecutionDispatch): + def __getattribute__(self, name): + if name in {"start_and_poll", "poll_bounded"}: + raise AssertionError("disabled fused dispatch operation was resolved") + return super().__getattribute__(name) + + +class _NonCallableExecutionDispatch(_TestExecutionDispatch): + start = 1 + + +class _NoneStartExecutionDispatch(_TestExecutionDispatch): + start = None + + +class _NoneFusedExecutionDispatch(_TestExecutionDispatch): + start_and_poll = None + + +class _UnretainedExecutionDispatch: + def __init__(self, slot: _FusedTestRequestSlot) -> None: + # None of these callable objects retains this dispatch container. + self.start = slot.start + self.start_and_poll = slot.start_and_poll + self.poll_bounded = slot.poll_bounded + self.poll_state = slot.poll_state + self.cancel = slot.cancel + self.recycle = slot.recycle + + +class _UnacknowledgedFusedRequestSlot(_FusedTestRequestSlot): + @property + def state(self) -> WorkState | None: + return None + + def poll_state(self) -> WorkState | None: + self.poll_count += 1 + return None + + def start_and_poll( + self, + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + self.fused_calls.append((max_polls, timeout_ns)) + self.start_count += 1 + raise KeyboardInterrupt + + +class _SlotReferenceBackend(_ReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.request_slots: list[_TestRequestSlot] = [] + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _TestRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.request_slots.append(slot) + + +class _FusedSlotReferenceBackend(_SlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + fused_request_slot_start_poll=True, + fused_request_slot_poll=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.request_slots.append(slot) + + +class _BlockingObservationSlotReferenceBackend(_SlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _BlockingObservationRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.request_slots.append(slot) + + +class _ExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.execution_dispatches: list[_TestExecutionDispatch] = [] + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + dispatch = _TestExecutionDispatch(slot) + try: + kwargs["adopt_slot"](slot, execution_dispatch=dispatch) + except BaseException: + return + self.request_slots.append(slot) + self.execution_dispatches.append(dispatch) + + +class _CapabilityGatedDispatchSlotBackend(_SlotReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.execution_dispatches: list[_CapabilityGatedExecutionDispatch] = [] + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + dispatch = _CapabilityGatedExecutionDispatch(slot) + try: + kwargs["adopt_slot"](slot, execution_dispatch=dispatch) + except BaseException: + return + self.request_slots.append(slot) + self.execution_dispatches.append(dispatch) + + +class _InvalidExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): + _DISPATCH_TYPE = _NonCallableExecutionDispatch + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + dispatch = self._DISPATCH_TYPE(slot) + try: + kwargs["adopt_slot"](slot, execution_dispatch=dispatch) + except BaseException: + return + raise AssertionError("non-callable execution dispatch was adopted") + + +class _NoneStartExecutionDispatchSlotBackend(_InvalidExecutionDispatchSlotBackend): + _DISPATCH_TYPE = _NoneStartExecutionDispatch + + +class _NoneFusedExecutionDispatchSlotBackend(_InvalidExecutionDispatchSlotBackend): + _DISPATCH_TYPE = _NoneFusedExecutionDispatch + + +class _UnretainedExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.execution_dispatch_ref = None + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + dispatch = _UnretainedExecutionDispatch(slot) + self.execution_dispatch_ref = weakref.ref(dispatch) + try: + kwargs["adopt_slot"](slot, execution_dispatch=dispatch) + except BaseException: + return + self.request_slots.append(slot) + + +class _StartFusedOnlySlotReferenceBackend(_FusedSlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + fused_request_slot_start_poll=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + +class _PollOnlySlotReferenceBackend(_FusedSlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + fused_request_slot_poll=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + +class _DeferredRecycleSlotReferenceBackend(_SlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + deferred_successful_request_slot_recycle=True, + ) + + +class _DeferredRecycleFusedSlotReferenceBackend(_FusedSlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + fused_request_slot_start_poll=True, + fused_request_slot_poll=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + deferred_successful_request_slot_recycle=True, + ) + + +class _NotificationSlotReferenceBackend(_FusedSlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + fused_request_slot_start_poll=True, + fused_request_slot_poll=True, + cuda_completion_event=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + deferred_successful_request_slot_recycle=True, + request_slot_notification_overrides=True, + ) + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _NotificationFusedTestRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.request_slots.append(slot) + + +class _MissingNotificationSlotBackend(_NotificationSlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _FusedTestRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + raise AssertionError("default notification dispatch was adopted") + + +class _UnacknowledgedFusedSlotBackend(_FusedSlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _UnacknowledgedFusedRequestSlot() + kwargs["adopt_slot"](slot) + self.request_slots.append(slot) + + +class _UnacknowledgedStartSlotBackend(_SlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _UnacknowledgedStartRequestSlot() + kwargs["adopt_slot"](slot) + self.request_slots.append(slot) + + +class _NormallyUnacknowledgedSlotBackend(_SlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _NormallyUnacknowledgedRequestSlot() + kwargs["adopt_slot"](slot) + self.request_slots.append(slot) + + +class _InterruptSlotCreationBackend(_SlotReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.created_slot: _TestRequestSlot | None = None + + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _TestRequestSlot() + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.created_slot = slot + raise KeyboardInterrupt + + +class _InterruptSlotCleanupBackend(_InterruptSlotCreationBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args + slot = _TestRequestSlot() + slot.fail_close_once = True + try: + kwargs["adopt_slot"](slot) + except BaseException: + return + self.created_slot = slot + raise KeyboardInterrupt + + +class _InterruptBeforeSlotAdoptionBackend(_SlotReferenceBackend): + def create_request_slot(self, *args, **kwargs) -> None: + del args, kwargs + raise KeyboardInterrupt + + +class _NoCompletionEventSlotBackend(_SlotReferenceBackend): + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + persistent_request_slots=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + +register_backend("test-pending", _PendingReferenceBackend, replace=True) +register_backend( + "test-counting-pending", _CountingPendingReferenceBackend, replace=True +) +register_backend( + "test-interrupting-failure", + _InterruptingFailureReferenceBackend, + replace=True, +) +register_backend("test-fused-selection", _FusedSelectionReferenceBackend, replace=True) +register_backend("test-background-only", _BackgroundOnlyReferenceBackend, replace=True) +register_backend("test-attached-only", _AttachedOnlyReferenceBackend, replace=True) +register_backend("test-bad-error", _BadErrorReferenceBackend, replace=True) +register_backend("test-interrupt", _InterruptReferenceBackend, replace=True) +register_backend( + "test-adopt-then-interrupt", _AdoptThenInterruptReferenceBackend, replace=True +) +register_backend( + "test-request-slot-no-completion-event", + _NoCompletionEventSlotBackend, + replace=True, +) +register_backend("test-regressing", _RegressingReferenceBackend, replace=True) +register_backend( + "test-transient-state-error", + _TransientStateErrorReferenceBackend, + replace=True, +) +register_backend( + "test-interrupt-after-close", _InterruptAfterCloseReferenceBackend, replace=True +) +register_backend("test-request-slot", _SlotReferenceBackend, replace=True) +register_backend( + "test-request-slot-blocking-observation", + _BlockingObservationSlotReferenceBackend, + replace=True, +) +register_backend("test-request-slot-fused", _FusedSlotReferenceBackend, replace=True) +register_backend( + "test-request-slot-execution-dispatch", + _ExecutionDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-execution-dispatch-gated", + _CapabilityGatedDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-execution-dispatch-invalid", + _InvalidExecutionDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-execution-dispatch-none-start", + _NoneStartExecutionDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-execution-dispatch-none-fused", + _NoneFusedExecutionDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-execution-dispatch-lifetime", + _UnretainedExecutionDispatchSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-start-fused-only", + _StartFusedOnlySlotReferenceBackend, + replace=True, +) +register_backend( + "test-request-slot-poll-only", + _PollOnlySlotReferenceBackend, + replace=True, +) +register_backend( + "test-request-slot-deferred-recycle", + _DeferredRecycleSlotReferenceBackend, + replace=True, +) +register_backend( + "test-request-slot-fused-deferred-recycle", + _DeferredRecycleFusedSlotReferenceBackend, + replace=True, +) +register_backend( + "test-request-slot-notification", + _NotificationSlotReferenceBackend, + replace=True, +) +register_backend( + "test-request-slot-notification-missing", + _MissingNotificationSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-fused-unacknowledged", + _UnacknowledgedFusedSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-unacknowledged-start", + _UnacknowledgedStartSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-no-ack", + _NormallyUnacknowledgedSlotBackend, + replace=True, +) +register_backend( + "test-request-slot-create-interrupt", + _InterruptSlotCreationBackend, + replace=True, +) +register_backend( + "test-request-slot-cleanup-interrupt", + _InterruptSlotCleanupBackend, + replace=True, +) +register_backend( + "test-request-slot-before-adoption-interrupt", + _InterruptBeforeSlotAdoptionBackend, + replace=True, +) + + +class EndpointTransferTest(TestCase): + def test_backend_registration_lost_return_retry_is_exactly_idempotent( + self, + ) -> None: + class Factory: + def __call__(self, *args, adopt_backend, **kwargs) -> None: + del args, adopt_backend, kwargs + + backend_name = "test-registration-lost-return" + factory = Factory() + register_backend(backend_name, factory, replace=True, factory_api_version=2) + + # Model a retry after Core committed the registry write but the caller + # was interrupted before observing the return. + register_backend(backend_name, factory, factory_api_version=2) + registration = _backend._get_backend_registration(backend_name) + self.assertIs(registration.factory, factory) + self.assertEqual(registration.factory_api_version, 2) + + # Idempotence is deliberately exact; neither a new callable nor a + # different contract can adopt an already-published name implicitly. + with self.assertRaises(BusyError): + register_backend(backend_name, factory, factory_api_version=1) + with self.assertRaises(BusyError): + register_backend(backend_name, Factory(), factory_api_version=2) + + def test_backend_factory_versions_are_explicit_and_never_introspected(self) -> None: + self.assertEqual(BACKEND_FACTORY_API_VERSION, 2) + for value, error_type in ( + (True, TypeError), + (2.0, TypeError), + (0, ValueError), + (3, ValueError), + ): + with self.subTest(value=value), self.assertRaises(error_type): + register_backend( + f"test-invalid-factory-version-{value!r}", + lambda *args: None, + factory_api_version=value, + ) + + events = [] + + class Factory: + @property + def __signature__(self): + raise AssertionError("factory signatures must not be inspected") + + def __call__(self, *args, adopt_backend, **kwargs) -> None: + del kwargs + backend = _ReferenceBackend._new_unactivated(*args) + events.append("pure") + self.assert_not_activated(backend) + adopt_backend(backend) + events.append("adopted") + backend._activate() + events.append("activated") + + @staticmethod + def assert_not_activated(backend) -> None: + if backend._activated: + raise AssertionError("native/global activation preceded adoption") + + backend_name = "test-factory-v2-no-introspection" + register_backend(backend_name, Factory(), replace=True, factory_api_version=2) + endpoint = Endpoint( + "factory-v2", backend=backend_name, endpoint_id="factory-v2-order" + ) + self.assertEqual(events, ["pure", "adopted", "activated"]) + self.assertTrue(endpoint._backend_receipt.published) + endpoint.close() + + # The legacy callable contract remains the default. + legacy_name = "test-factory-v1-compatible" + register_backend(legacy_name, _ReferenceBackend, replace=True) + legacy = Endpoint("factory-v1", backend=legacy_name) + self.assertEqual( + _backend._get_backend_registration(legacy_name).factory_api_version, 1 + ) + legacy.close() + + def test_backend_factory_v2_rejects_invalid_adoption_contracts(self) -> None: + retained_callbacks = [] + + for behavior in ("missing", "duplicate", "wrong", "non-none"): + instances = [] + + def factory(*args, adopt_backend, _behavior=behavior, **kwargs): + del kwargs + if _behavior == "missing": + retained_callbacks.append(adopt_backend) + return None + if _behavior == "wrong": + adopt_backend(object()) + return None + backend = _ReferenceBackend._new_unactivated(*args) + instances.append(backend) + adopt_backend(backend) + if _behavior == "duplicate": + adopt_backend(backend) + backend._activate() + return object() if _behavior == "non-none" else None + + backend_name = f"test-factory-v2-invalid-{behavior}" + register_backend(backend_name, factory, replace=True, factory_api_version=2) + with ( + self.subTest(behavior=behavior), + self.assertRaises((BackendFailureError, TypeError)), + ): + Endpoint("invalid-v2", backend=backend_name) + for instance in instances: + self.assertTrue(instance._closed) + + late_backend = _ReferenceBackend( + "late-adoption", + "late-adoption", + "late-incarnation", + _types.ProgressMode.MANUAL, + ThreadMode.SERIALIZED, + {}, + ) + with self.assertRaisesRegex(BackendFailureError, "after its call returned"): + retained_callbacks[0](late_backend) + late_backend.close() + + def test_backend_factory_v1_wrong_type_is_foreign_and_never_touched(self) -> None: + calls = [] + + class ForeignResult: + def close(self) -> None: + calls.append("close") + + candidate = ForeignResult() + backend_name = "test-factory-v1-foreign-result" + register_backend( + backend_name, + lambda *args, **kwargs: candidate, + replace=True, + factory_api_version=1, + ) + + with self.assertRaisesRegex(TypeError, "must return a Backend instance"): + Endpoint("foreign-result", backend=backend_name) + self.assertEqual(calls, []) + + def test_backend_factory_v2_post_adoption_failure_exposes_recovery(self) -> None: + instances = [] + + class PendingCloseBackend(_ReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def close(self) -> None: + if not hasattr(self, "allow_close"): + self.allow_close = False + if not self.allow_close: + raise BusyError("injected activation cleanup hold") + super().close() + + def factory(*args, adopt_backend, **kwargs) -> None: + del kwargs + backend = PendingCloseBackend._new_unactivated(*args) + instances.append(backend) + adopt_backend(backend) + backend._activate() + raise KeyboardInterrupt("post-adoption activation cut") + + backend_name = "test-factory-v2-post-adoption-recovery" + register_backend(backend_name, factory, replace=True, factory_api_version=2) + with self.assertRaisesRegex( + BackendFailureError, "recovery_endpoint" + ) as context: + Endpoint("interrupted-v2", backend=backend_name) + recovery = context.exception.recovery_endpoint + self.assertIs(recovery._backend, instances[0]) + self.assertFalse(recovery._backend_receipt.published) + self.assertTrue(recovery._creation_failed) + instances[0].allow_close = True + recovery.close() + self.assertTrue(recovery.closed) + + def test_backend_finalizer_orphans_are_isolated_by_exact_identity(self) -> None: + blocked_ids = {"published-orphan"} + instances = [] + + class OrphanBackend(_ReferenceBackend): + def close(self) -> None: + if self._endpoint_id in blocked_ids: + raise BusyError("injected live backend child") + super().close() + + def factory(*args, adopt_backend, **kwargs) -> None: + del kwargs + backend = OrphanBackend._new_unactivated(*args) + instances.append(backend) + adopt_backend(backend) + backend._activate() + + backend_name = "test-factory-v2-published-orphan" + register_backend(backend_name, factory, replace=True, factory_api_version=2) + abandoned = Endpoint( + "abandoned", backend=backend_name, endpoint_id="published-orphan" + ) + abandoned_ref = weakref.ref(abandoned) + del abandoned + gc.collect() + self.assertIsNone(abandoned_ref()) + self.assertIn( + (backend_name, "published-orphan"), + _api._ORPHANED_BACKEND_RECEIPTS, + ) + + # Auditing the failed orphan does not gate an unrelated endpoint id. + unrelated = Endpoint( + "unrelated", backend=backend_name, endpoint_id="fresh-identity" + ) + unrelated.close() + call_count = len(instances) + with self.assertRaisesRegex( + BackendFailureError, "abandoned published endpoint" + ): + Endpoint("collision", backend=backend_name, endpoint_id="published-orphan") + self.assertEqual(len(instances), call_count) + + blocked_ids.clear() + replacement = Endpoint( + "replacement", backend=backend_name, endpoint_id="published-orphan" + ) + replacement.close() + self.assertNotIn( + (backend_name, "published-orphan"), + _api._ORPHANED_BACKEND_RECEIPTS, + ) + + def test_unpublished_factory_orphan_has_distinct_exact_id_diagnostic(self) -> None: + block_cleanup = True + fail_factory = True + + class OrphanBackend(_ReferenceBackend): + def close(self) -> None: + if block_cleanup: + raise BusyError("injected construction cleanup hold") + super().close() + + def factory(*args, adopt_backend, **kwargs) -> None: + del kwargs + backend = OrphanBackend._new_unactivated(*args) + adopt_backend(backend) + backend._activate() + if fail_factory: + raise RuntimeError("injected post-activation failure") + + backend_name = "test-factory-v2-unpublished-orphan" + register_backend(backend_name, factory, replace=True, factory_api_version=2) + with self.assertRaises(BackendFailureError) as context: + Endpoint( + "unpublished", backend=backend_name, endpoint_id="unpublished-orphan" + ) + recovery = context.exception.recovery_endpoint + recovery_ref = weakref.ref(recovery) + # The test harness retains the raised exception until the test returns; + # sever its documented recovery edge to model an application dropping + # both objects and allow the Endpoint weak finalizer to run now. + context.exception.recovery_endpoint = None + context.exception.__traceback__ = None + context.exception.__cause__ = None + del recovery + del context + gc.collect() + self.assertIsNone(recovery_ref()) + + fail_factory = False + with self.assertRaisesRegex( + BackendFailureError, "unpublished endpoint construction" + ): + Endpoint( + "collision", backend=backend_name, endpoint_id="unpublished-orphan" + ) + block_cleanup = False + replacement = Endpoint( + "replacement", backend=backend_name, endpoint_id="unpublished-orphan" + ) + replacement.close() + + def test_endpoint_close_cut_after_receipt_commit_is_retryable(self) -> None: + endpoint = Endpoint("receipt-close-cut") + original_detach = Endpoint._detach_backend_finalizer + interrupt = True + + def detach_with_cut(owner) -> None: + nonlocal interrupt + if interrupt: + interrupt = False + raise KeyboardInterrupt("after receipt close commit") + original_detach(owner) + + with mock.patch.object(Endpoint, "_detach_backend_finalizer", detach_with_cut): + with self.assertRaisesRegex( + KeyboardInterrupt, "after receipt close commit" + ): + endpoint.close() + self.assertTrue(endpoint._backend_receipt.closed) + self.assertFalse(endpoint.closed) + self.assertTrue(endpoint._backend_finalizer.alive) + endpoint.close() + self.assertTrue(endpoint.closed) + self.assertFalse(endpoint._backend_finalizer.alive) + + def test_endpoint_close_does_not_implicitly_progress(self) -> None: + class NoCloseProgressBackend(_ReferenceBackend): + def progress(self) -> None: + raise AssertionError("close must not inject provider progress") + + backend_name = "test-close-without-progress" + register_backend(backend_name, NoCloseProgressBackend, replace=True) + endpoint = Endpoint("no-close-progress", backend=backend_name) + endpoint.close() + self.assertTrue(endpoint.closed) + + def test_backend_receipt_empty_slots_preserve_unslotted_subclasses(self) -> None: + self.assertEqual(BackendWork.__slots__, ()) + self.assertEqual(BackendRequestSlot.__slots__, ()) + + work = _PendingWork() + slot = _TestRequestSlot() + work.extension_value = "work" + slot.extension_value = "slot" + + self.assertEqual(work.__dict__["extension_value"], "work") + self.assertEqual(slot.__dict__["extension_value"], "slot") + + def test_endpoint_close_retries_without_progress_after_interruption(self) -> None: + endpoint = Endpoint("endpoint", backend="test-interrupt-after-close") + + with self.assertRaises(KeyboardInterrupt): + endpoint.close() + self.assertFalse(endpoint.closed) + self.assertTrue(endpoint._closing) + + endpoint.close() + self.assertTrue(endpoint.closed) + + def test_registration_factory_interruption_rolls_back_adopted_handle(self) -> None: + endpoint = Endpoint("endpoint") + original_register = endpoint._backend.register + + def register_then_interrupt(*args, **kwargs) -> None: + original_register(*args, **kwargs) + raise KeyboardInterrupt + + endpoint._backend.register = register_then_interrupt + with self.assertRaises(KeyboardInterrupt): + endpoint.register(bytearray(4), name="buffer") + + self.assertEqual(endpoint.open_registrations, ()) + self.assertFalse(endpoint._backend._registrations) + endpoint.close() + + def test_failed_registration_rollback_is_publicly_recoverable(self) -> None: + endpoint = Endpoint("endpoint") + original_register = endpoint._backend.register + original_deregister = endpoint._backend.deregister + fail_cleanup = True + + def register_then_interrupt(*args, **kwargs) -> None: + original_register(*args, **kwargs) + raise KeyboardInterrupt + + def deregister_once(handle) -> None: + nonlocal fail_cleanup + if fail_cleanup: + fail_cleanup = False + raise RuntimeError("injected cleanup failure") + original_deregister(handle) + + endpoint._backend.register = register_then_interrupt + endpoint._backend.deregister = deregister_once + with self.assertRaisesRegex(BackendFailureError, "cleanup is pending"): + endpoint.register(bytearray(4), name="buffer") + + (registration,) = endpoint.open_registrations + with self.assertRaisesRegex(BackendFailureError, "cleanup only"): + registration.region() + registration.close() + self.assertEqual(endpoint.open_registrations, ()) + endpoint.close() + + def test_peer_and_plan_factory_interruptions_roll_back(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(4), name="local") + remote = destination.register(bytearray(4), name="remote") + metadata = bytes(destination.export_metadata()) + original_import = source._backend.import_peer + + def import_then_interrupt(*args, **kwargs) -> None: + original_import(*args, **kwargs) + raise KeyboardInterrupt + + source._backend.import_peer = import_then_interrupt + with self.assertRaises(KeyboardInterrupt): + source.import_peer(metadata) + self.assertEqual(source.open_peers, ()) + self.assertFalse(source._backend._peers) + + source._backend.import_peer = original_import + peer = source.import_peer(metadata) + original_prepare = source._backend.prepare + + def prepare_then_interrupt(*args, **kwargs) -> None: + original_prepare(*args, **kwargs) + raise KeyboardInterrupt + + source._backend.prepare = prepare_then_interrupt + with self.assertRaises(KeyboardInterrupt): + source.prepare(local=[local.region()], remote=[peer.region("remote")]) + self.assertEqual(source.open_plans, ()) + self.assertFalse(source._backend._plans) + + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_reference_write_and_metadata_round_trip(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"abcdefgh") + destination_buffer = bytearray(8) + source_registration = source.register(source_buffer, name="blocks") + destination_registration = destination.register( + destination_buffer, name="blocks" + ) + + encoded = bytes(destination.export_metadata()) + self.assertEqual(bytes(EndpointMetadata.from_bytes(encoded)), encoded) + peer = source.import_peer(encoded) + plan = source.prepare( + local=[source_registration.region()], + remote=[peer.region("blocks")], + ) + work = source.write(plan) + self.assertTrue(work.test()) + self.assertEqual(work.state, WorkState.COMPLETED) + self.assertEqual(destination_buffer, source_buffer) + + work.close() + plan.close() + peer.close() + source_registration.close() + destination_registration.close() + source.close() + destination.close() + + def test_reference_read(self) -> None: + sender = Endpoint("sender") + receiver = Endpoint("receiver") + source_buffer = bytearray(b"remote") + destination_buffer = bytearray(6) + source_registration = sender.register(source_buffer, name="weights") + destination_registration = receiver.register(destination_buffer, name="weights") + peer = receiver.import_peer(bytes(sender.export_metadata())) + plan = receiver.prepare( + local=[destination_registration.region()], + remote=[peer.region("weights")], + ) + + work = receiver.read(plan) + work.wait() + self.assertEqual(destination_buffer, source_buffer) + + work.close() + plan.close() + peer.close() + destination_registration.close() + source_registration.close() + receiver.close() + sender.close() + + def test_bound_transfer_full_reuse_and_provider_tokens(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"data") + destination_buffer = bytearray(4) + local = source.register(source_buffer, name="buffer") + remote = destination.register(destination_buffer, name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + tokens: list[object | None] = [] + original_submit = source._backend.submit + + def capture_submit(*args, **kwargs) -> None: + tokens.append(kwargs["reuse_token"]) + original_submit(*args, **kwargs) + + source._backend.submit = capture_submit + first = plan.bind("write") + self.assertIsInstance(first, BoundTransfer) + self.assertEqual(first.op, TransferOp.WRITE) + self.assertEqual(source.open_bound_transfers, (first,)) + + for value in (b"one!", b"two!"): + source_buffer[:] = value + work = first.submit() + work.wait() + self.assertEqual(destination_buffer, value) + work.close() + + second = plan.bind(TransferOp.WRITE) + work = second.submit() + work.close() + reader = plan.bind(TransferOp.READ) + destination_buffer[:] = b"back" + work = reader.submit() + work.wait() + self.assertEqual(source_buffer, b"back") + work.close() + + convenience_start = len(tokens) + for _ in range(2): + work = source.write(plan) + work.close() + + self.assertIs(tokens[0], tokens[1]) + self.assertIsNot(tokens[0], tokens[2]) + self.assertIsNot(tokens[0], tokens[3]) + self.assertIsNot(tokens[2], tokens[3]) + self.assertIs(tokens[convenience_start], tokens[convenience_start + 1]) + self.assertIsNot(tokens[0], tokens[convenience_start]) + + reader.close() + second.close() + first.close() + self.assertEqual(source.open_bound_transfers, ()) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_bound_transfer_canonicalizes_indexed_selection_and_notification( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"a0b0c0d0") + destination_buffer = bytearray(b"--------") + local = source.register(source_buffer, name="source") + remote = destination.register(destination_buffer, name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=2, count=4)], + remote=[peer.region("destination", 0, 1, stride=2, count=4)], + indexed=True, + ) + local_indices = [3, 1] + remote_indices = [0, 2] + notification = bytearray(b"ready") + bound = plan.bind( + TransferOp.WRITE, + local_indices, + remote_indices, + notification, + ) + local_indices[:] = [0] + remote_indices[:] = [1] + notification[:] = b"later" + + work = bound.submit() + work.wait() + self.assertEqual(destination_buffer, b"d---b---") + self.assertEqual( + [item.payload for item in destination.poll_notifications()], + [b"ready"], + ) + + work.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_bound_transfer_enforces_work_plan_endpoint_close_order(self) -> None: + source = Endpoint("source", backend="test-pending") + destination = Endpoint("destination", backend="test-pending") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaisesRegex(BusyError, "bound transfers"): + plan.close() + with self.assertRaises(BusyError): + source.close() + work = bound.submit() + with self.assertRaisesRegex(BusyError, "open work"): + bound.close() + with self.assertRaisesRegex(BusyError, "open work"): + plan.close() + self.assertIn(work, source.open_works) + self.assertIn(bound, source.open_bound_transfers) + + self.assertTrue(work.cancel()) + work.close() + bound.close() + self.assertTrue(bound.closed) + with self.assertRaises(ClosedError): + bound.submit() + self.assertEqual(source.open_bound_transfers, ()) + + plan.close() + with self.assertRaises(ClosedError): + plan.bind(TransferOp.WRITE) + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_bound_transfer_preserves_adoption_interruption_recovery(self) -> None: + source = Endpoint("source", backend="test-adopt-then-interrupt") + destination = Endpoint("destination", backend="test-adopt-then-interrupt") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + work = bound.submit() + self.assertEqual(source.open_works, (work,)) + work.cancel() + self.assertEqual(work.state, WorkState.FAILED) + self.assertIsInstance(work.error, KeyboardInterrupt) + work.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_bound_transfer_rolls_back_interruption_before_adoption(self) -> None: + source = Endpoint("source", backend="test-interrupt") + destination = Endpoint("destination", backend="test-interrupt") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaises(KeyboardInterrupt): + bound.submit() + self.assertEqual(source.open_works, ()) + self.assertFalse(bound._works) + self.assertFalse(plan._submitting) + self.assertFalse(plan._works) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_bound_transfer_close_retry_repairs_post_commit_interruption(self) -> None: + class InterruptingBoundTransfer(BoundTransfer): + __slots__ = ("interrupt_close_commit",) + + def __setattr__(self, name: str, value: object) -> None: + object.__setattr__(self, name, value) + if ( + name == "_closed" + and value is True + and getattr(self, "interrupt_close_commit", False) + ): + object.__setattr__(self, "interrupt_close_commit", False) + raise KeyboardInterrupt + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + with mock.patch.object(_api, "BoundTransfer", InterruptingBoundTransfer): + bound = plan.bind(TransferOp.WRITE) + bound.interrupt_close_commit = True + + with self.assertRaises(KeyboardInterrupt): + bound.close() + self.assertTrue(bound.closed) + # The parent blockers were removed before the interrupted commit, but + # the retained parent reference lets the same-operation retry finish. + self.assertNotIn(bound, plan._bindings) + self.assertNotIn(bound, source.open_bound_transfers) + self.assertIs(bound._plan, plan) + + bound.close() + self.assertIsNone(bound._plan) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_resource_close_retries_finish_post_commit_epilogues(self) -> None: + class InterruptOnCloseCommit: + __slots__ = () + + def __setattr__(self, name: str, value: object) -> None: + object.__setattr__(self, name, value) + if ( + name == "_closed" + and value is True + and getattr(self, "interrupt_close_commit", False) + ): + object.__setattr__(self, "interrupt_close_commit", False) + raise KeyboardInterrupt + + class InterruptingRegistration(InterruptOnCloseCommit, _api.Registration): + __slots__ = ("interrupt_close_commit",) + + class InterruptingPeer(InterruptOnCloseCommit, _api.Peer): + __slots__ = ("interrupt_close_commit",) + + class InterruptingPlan(InterruptOnCloseCommit, _api.TransferPlan): + __slots__ = ("interrupt_close_commit",) + + class InterruptingWork(InterruptOnCloseCommit, _api.Work): + __slots__ = ("interrupt_close_commit",) + + source = Endpoint("source") + destination = Endpoint("destination") + with mock.patch.object(_api, "Registration", InterruptingRegistration): + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + with mock.patch.object(_api, "Peer", InterruptingPeer): + peer = source.import_peer(bytes(destination.export_metadata())) + with mock.patch.object(_api, "TransferPlan", InterruptingPlan): + plan = source.prepare( + local=[local.region()], remote=[peer.region("buffer")] + ) + with mock.patch.object(_api, "Work", InterruptingWork): + work = source.write(plan) + work.wait() + + for resource in (work, plan, peer, local): + resource.interrupt_close_commit = True + with self.assertRaises(KeyboardInterrupt): + resource.close() + self.assertTrue(resource.closed) + + if resource is work: + self.assertIn(work, source.open_works) + elif resource is plan: + self.assertIn(plan, source.open_plans) + self.assertTrue(plan._origins) + elif resource is peer: + self.assertIn(peer, source.open_peers) + self.assertIsNotNone(peer._handle) + else: + self.assertIn(local, source.open_registrations) + self.assertIsNotNone(local._memory.owner) + + resource.close() + + self.assertEqual(source.open_works, ()) + self.assertEqual(source.open_plans, ()) + self.assertEqual(source.open_peers, ()) + self.assertEqual(source.open_registrations, ()) + remote.close() + source.close() + destination.close() + + def test_bound_host_wait_establishes_ordering_on_every_submit(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local.region() + local_region._device = "cuda:0" + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + source._capabilities = Capabilities( + read=True, + write=True, + manual_progress=True, + background_progress=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + ) + actions: list[str] = [] + + class FakeEvent: + device = torch.device("cuda:0") + + def synchronize(self) -> None: + actions.append("wait") + + with ( + mock.patch.object(torch.cuda, "Event", FakeEvent), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=False + ), + ): + bound = plan.bind( + TransferOp.WRITE, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(), + ) + self.assertEqual(actions, []) + for _ in range(2): + work = bound.submit() + work.close() + self.assertEqual(actions, ["wait", "wait"]) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_backend_request_slot_default_start_and_poll_is_bounded(self) -> None: + class CountingRequestSlot(_TestRequestSlot): + def __init__(self) -> None: + super().__init__() + self.next_state = WorkState.RUNNING + self.poll_count = 0 + + def poll_state(self) -> WorkState: + self.poll_count += 1 + return self.state + + backend_slot = CountingRequestSlot() + self.assertEqual( + backend_slot.start_and_poll(max_polls=3, timeout_ns=None), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.start_count, 1) + self.assertEqual(backend_slot.poll_count, 3) + + backend_slot = CountingRequestSlot() + self.assertEqual( + backend_slot.start_and_poll(max_polls=64, timeout_ns=0), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.poll_count, 1) + backend_slot = CountingRequestSlot() + self.assertEqual( + backend_slot.start_and_poll(max_polls=0, timeout_ns=None), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.start_count, 1) + self.assertEqual(backend_slot.poll_count, 0) + with self.assertRaisesRegex(ValueError, "max_polls"): + backend_slot.start_and_poll(max_polls=-1, timeout_ns=None) + with self.assertRaisesRegex(ValueError, "max_polls"): + backend_slot.start_and_poll(max_polls=1 << 63, timeout_ns=None) + with self.assertRaisesRegex(TypeError, "max_polls"): + backend_slot.start_and_poll(max_polls=True, timeout_ns=None) + with self.assertRaisesRegex(ValueError, "timeout_ns"): + backend_slot.start_and_poll(max_polls=1, timeout_ns=-1) + with self.assertRaisesRegex(ValueError, "timeout_ns"): + backend_slot.start_and_poll(max_polls=1, timeout_ns=1 << 63) + with self.assertRaisesRegex(TypeError, "timeout_ns"): + backend_slot.start_and_poll(max_polls=1, timeout_ns=True) + + backend_slot = CountingRequestSlot() + backend_slot.start = mock.Mock(return_value=object()) + with self.assertRaisesRegex(TypeError, "start must return None"): + backend_slot.start_and_poll(max_polls=1, timeout_ns=None) + + def test_backend_request_slot_default_poll_bounded_contract(self) -> None: + class CountingRequestSlot(_TestRequestSlot): + def __init__(self) -> None: + super().__init__() + self._state = WorkState.RUNNING + self.poll_count = 0 + + def poll_state(self) -> WorkState: + self.poll_count += 1 + assert isinstance(self._state, WorkState) + return self._state + + backend_slot = CountingRequestSlot() + self.assertEqual( + backend_slot.poll_bounded( + max_polls=3, timeout_ns=None, timeout_check_interval=32 + ), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.poll_count, 3) + + backend_slot = CountingRequestSlot() + with mock.patch( + "torch.distributed._transfer._backend.time.monotonic_ns", + side_effect=[10, 10], + ): + self.assertEqual( + backend_slot.poll_bounded( + max_polls=64, timeout_ns=0, timeout_check_interval=32 + ), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.poll_count, 1) + + backend_slot = CountingRequestSlot() + with mock.patch( + "torch.distributed._transfer._backend.time.monotonic_ns", + side_effect=[10, 10, 12], + ): + self.assertEqual( + backend_slot.poll_bounded( + max_polls=64, timeout_ns=1, timeout_check_interval=1 + ), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.poll_count, 2) + + backend_slot = CountingRequestSlot() + with mock.patch( + "torch.distributed._transfer._backend.time.monotonic_ns", + side_effect=[10, 10, 12], + ): + self.assertEqual( + backend_slot.poll_bounded( + max_polls=64, timeout_ns=1, timeout_check_interval=32 + ), + WorkState.RUNNING, + ) + self.assertEqual(backend_slot.poll_count, 32) + + for max_polls, timeout_ns, timeout_check_interval, error_type in ( + (0, None, 32, ValueError), + (True, None, 32, TypeError), + (1, -1, 32, ValueError), + (1 << 63, None, 32, ValueError), + (1, 1 << 63, 32, ValueError), + (1, True, 32, TypeError), + (1, None, 0, ValueError), + (1, None, 1 << 63, ValueError), + (1, None, True, TypeError), + ): + with self.assertRaises(error_type): + backend_slot.poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + + def test_transfer_slot_capability_is_validated_and_opt_in(self) -> None: + with self.assertRaisesRegex(TypeError, "persistent_request_slots"): + Capabilities(persistent_request_slots=1) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "fused_request_slot_start_poll"): + Capabilities( + persistent_request_slots=True, + fused_request_slot_start_poll=1, # type: ignore[arg-type] + ) + with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): + Capabilities(fused_request_slot_start_poll=True) + with self.assertRaisesRegex(TypeError, "fused_request_slot_poll"): + Capabilities( + persistent_request_slots=True, + fused_request_slot_poll=1, # type: ignore[arg-type] + ) + with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): + Capabilities(fused_request_slot_poll=True) + with self.assertRaisesRegex( + TypeError, "deferred_successful_request_slot_recycle" + ): + Capabilities( + persistent_request_slots=True, + deferred_successful_request_slot_recycle=1, # type: ignore[arg-type] + ) + with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): + Capabilities(deferred_successful_request_slot_recycle=True) + with self.assertRaisesRegex(TypeError, "request_slot_notification_overrides"): + Capabilities( + persistent_request_slots=True, + attached_notifications=True, + request_slot_notification_overrides=1, # type: ignore[arg-type] + ) + with self.assertRaisesRegex(TypeError, "fused_prevalidated_submit_poll"): + Capabilities(fused_prevalidated_submit_poll=1) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "fused_work_poll"): + Capabilities(fused_work_poll=1) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "grouped_notification_receive"): + Capabilities(grouped_notification_receive=1) # type: ignore[arg-type] + with self.assertRaisesRegex(ValueError, "requires notifications"): + Capabilities(grouped_notification_receive=True) + with self.assertRaisesRegex(TypeError, "synchronous_notification_send"): + Capabilities(synchronous_notification_send=1) # type: ignore[arg-type] + with self.assertRaisesRegex(ValueError, "requires standalone_notifications"): + Capabilities(synchronous_notification_send=True) + synchronous_notifications = Capabilities( + standalone_notifications=True, + standalone_notification_completion=NotificationCompletion.LOCAL_ACCEPTED, + synchronous_notification_send=True, + ) + self.assertTrue(synchronous_notifications.synchronous_notification_send) + self.assertTrue( + Capabilities( + attached_notifications=True, + grouped_notification_receive=True, + ).grouped_notification_receive + ) + with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): + Capabilities( + attached_notifications=True, + request_slot_notification_overrides=True, + ) + with self.assertRaisesRegex(ValueError, "requires attached_notifications"): + Capabilities( + persistent_request_slots=True, + request_slot_notification_overrides=True, + ) + poll_only = Capabilities( + persistent_request_slots=True, + fused_request_slot_poll=True, + ) + self.assertTrue(poll_only.fused_request_slot_poll) + self.assertFalse(poll_only.fused_request_slot_start_poll) + self.assertFalse(poll_only.deferred_successful_request_slot_recycle) + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + self.assertFalse(source.capabilities.persistent_request_slots) + self.assertFalse(source.capabilities.fused_request_slot_start_poll) + self.assertFalse(source.capabilities.fused_request_slot_poll) + self.assertFalse(source.capabilities.deferred_successful_request_slot_recycle) + self.assertFalse(source.capabilities.request_slot_notification_overrides) + self.assertFalse(source.capabilities.fused_prevalidated_submit_poll) + self.assertFalse(source.capabilities.fused_work_poll) + with self.assertRaisesRegex(UnsupportedError, "persistent request slots"): + bound.create_slot() + self.assertEqual(source.open_transfer_slots, ()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_notification_slot_is_explicit_reusable_and_work_free( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-notification") + destination = Endpoint("destination", backend="test-request-slot-notification") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaisesRegex(TypeError, "must be a bool"): + bound.create_slot(request_local_notifications=1) # type: ignore[arg-type] + + fixed_slot = bound.create_slot() + self.assertIs(type(fixed_slot), TransferSlot) + fixed_slot.close() + + slot = bound.create_slot(request_local_notifications=True) + self.assertIs(type(slot), RequestLocalNotificationSlot) + self.assertIsInstance(slot, TransferSlot) + self.assertEqual(source.open_works, ()) + backend_slot = source._backend.request_slots[-1] + + with self.assertRaisesRegex(TypeError, "notification"): + slot.start() # type: ignore[call-arg] + self.assertEqual(slot.generation, 0) + self.assertEqual(backend_slot.start_count, 0) + + released_view = memoryview(bytearray(b"released")) + released_view.release() + for invalid, error_type, message in ( + (b"", ValueError, "must be non-empty"), + (b"x" * (_api._MAX_NOTIFICATION_BYTES + 1), ValueError, "size limit"), + ("wrong", TypeError, "bytes-like"), + (released_view, ValueError, "unavailable"), + ): + with self.assertRaisesRegex(error_type, message): + slot.start(notification=invalid) # type: ignore[arg-type] + self.assertEqual(slot.generation, 0) + self.assertEqual(backend_slot.start_count, 0) + + mutable = bytearray(b"alpha") + slot.start(notification=mutable) + mutable[:] = b"other" + self.assertEqual(backend_slot.notification_starts, [b"alpha"]) + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + + view_owner = bytearray(b"view") + slot.start(notification=memoryview(view_owner)) + view_owner[:] = b"xxxx" + self.assertEqual(backend_slot.notification_starts[-1], b"view") + slot.recycle() + + generation = slot.execute(notification=b"b\x00\xff") + self.assertEqual(generation, 3) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual( + backend_slot.notification_fused_calls, + [(b"b\x00\xff", 64, None)], + ) + + self.assertEqual(slot.execute(notification=None), 4) + self.assertEqual( + backend_slot.notification_fused_calls[-1], + (None, 64, None), + ) + self.assertEqual(source.open_works, ()) + + prior_calls = len(backend_slot.notification_fused_calls) + with self.assertRaisesRegex(ValueError, "must be non-empty"): + slot.execute(notification=b"") + self.assertEqual(slot.generation, 4) + self.assertEqual(len(backend_slot.notification_fused_calls), prior_calls) + + tagged = plan.bind(TransferOp.WRITE, notification=b"fixed") + with self.assertRaisesRegex(ValueError, "notification=None"): + tagged.create_slot(request_local_notifications=True) + tagged.close() + + second = bound.create_slot(request_local_notifications=True) + slot.start(notification=b"one") + second.start(notification=b"two") + self.assertEqual(slot.state, WorkState.COMPLETED) + self.assertEqual(second.state, WorkState.COMPLETED) + slot.recycle() + second.recycle() + second.close() + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_notification_slots_cover_indexed_write_and_read( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-notification") + destination = Endpoint("destination", backend="test-request-slot-notification") + local = source.register(bytearray(b"abcdefgh"), name="buffer") + remote = destination.register(bytearray(8), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 4, stride=4, count=2)], + remote=[peer.region("buffer", 0, 4, stride=4, count=2)], + indexed=True, + ) + + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[1], + remote_indices=[0], + ) + write_bound = selection.bind() + write_slot = write_bound.create_slot(request_local_notifications=True) + self.assertEqual(write_slot.execute(notification=b"indexed"), 1) + + read_bound = plan.bind( + TransferOp.READ, + local_indices=[0], + remote_indices=[1], + ) + read_slot = read_bound.create_slot(request_local_notifications=True) + self.assertEqual(read_slot.execute(notification=b"read"), 1) + + read_slot.close() + read_bound.close() + write_slot.close() + write_bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_notification_slot_is_capability_and_spi_gated( + self, + ) -> None: + for backend, message in ( + ("test-request-slot", "does not support request-local"), + ("test-request-slot-notification-missing", "without adoption"), + ): + source = Endpoint("source", backend=backend) + destination = Endpoint("destination", backend=backend) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region()], remote=[peer.region("buffer")] + ) + bound = plan.bind(TransferOp.WRITE) + with self.assertRaisesRegex( + (UnsupportedError, BackendFailureError), message + ): + bound.create_slot(request_local_notifications=True) + self.assertEqual(source.open_transfer_slots, ()) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_reuses_one_public_resource_and_counts_generations( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + self.assertIsInstance(slot, TransferSlot) + self.assertEqual(source.open_transfer_slots, (slot,)) + self.assertEqual(source.open_works, ()) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot.generation, 0) + + self.assertIsNone(slot.start()) + self.assertTrue(slot.active) + self.assertEqual(slot.generation, 1) + self.assertEqual(slot.state, WorkState.COMPLETED) + self.assertIs(slot.wait(), slot) + backend_slot = source._backend.request_slots[0] + self.assertIs(slot.completion_event, backend_slot.completion) + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.start() + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.close() + + slot.recycle() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot.generation, 1) + self.assertIsNone(slot.start()) + self.assertEqual(slot.generation, 2) + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + self.assertEqual(backend_slot.start_count, 2) + self.assertEqual(backend_slot.recycle_count, 2) + self.assertEqual(source.open_works, ()) + + slot.close() + self.assertTrue(slot.closed) + self.assertEqual(source.open_transfer_slots, ()) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execution_dispatch_is_cached_and_wrapper_owned( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-execution-dispatch") + destination = Endpoint( + "destination", backend="test-request-slot-execution-dispatch" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + dispatch = source._backend.execution_dispatches[0] + self.assertIsNone(slot._fast_execute) + + self.assertEqual( + dispatch.attribute_reads, + {name: 1 for name in _TestExecutionDispatch._OPERATION_NAMES}, + ) + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + self.assertIs(slot.start_and_wait(), slot) + self.assertIs(slot.completion_event, backend_slot.completion) + slot.recycle() + + backend_slot.next_state = WorkState.RUNNING + slot.start() + self.assertEqual(slot.state, WorkState.RUNNING) + self.assertTrue(slot.cancel()) + self.assertEqual(slot.state, WorkState.CANCELLED) + slot.recycle() + + failure = RuntimeError("wrapper-owned diagnostic") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + slot.start() + self.assertIs(slot.error, failure) + slot.recycle() + self.assertEqual( + dispatch.attribute_reads, + {name: 1 for name in _TestExecutionDispatch._OPERATION_NAMES}, + ) + + slot.close() + self.assertEqual(backend_slot.close_count, 1) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execution_dispatch_respects_capability_gating( + self, + ) -> None: + source = Endpoint( + "source", backend="test-request-slot-execution-dispatch-gated" + ) + destination = Endpoint( + "destination", backend="test-request-slot-execution-dispatch-gated" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + dispatch = source._backend.execution_dispatches[0] + + self.assertEqual(dispatch.attribute_reads["start"], 1) + self.assertEqual(dispatch.attribute_reads["poll_state"], 1) + self.assertEqual(dispatch.attribute_reads["cancel"], 1) + self.assertEqual(dispatch.attribute_reads["recycle"], 1) + self.assertEqual(dispatch.attribute_reads["start_and_poll"], 0) + self.assertEqual(dispatch.attribute_reads["poll_bounded"], 0) + slot.start() + slot.recycle() + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_strongly_retains_execution_dispatch_until_close( + self, + ) -> None: + source = Endpoint( + "source", backend="test-request-slot-execution-dispatch-lifetime" + ) + destination = Endpoint( + "destination", backend="test-request-slot-execution-dispatch-lifetime" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + dispatch_ref = source._backend.execution_dispatch_ref + + gc.collect() + self.assertIsNotNone(dispatch_ref()) + self.assertIs(slot._backend_execution_dispatch, dispatch_ref()) + slot.close() + gc.collect() + self.assertIsNone(dispatch_ref()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execution_dispatch_rejects_noncallable_operation( + self, + ) -> None: + for backend in ( + "test-request-slot-execution-dispatch-invalid", + "test-request-slot-execution-dispatch-none-start", + "test-request-slot-execution-dispatch-none-fused", + ): + with self.subTest(backend=backend): + source = Endpoint("source", backend=backend) + destination = Endpoint("destination", backend=backend) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region()], remote=[peer.region("buffer")] + ) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaisesRegex(BackendFailureError, "without adoption"): + bound.create_slot() + self.assertEqual(source.open_transfer_slots, ()) + self.assertEqual(bound._slots, set()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_pending_cancel_failure_and_recycle_contract(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + with self.assertRaisesRegex(BusyError, "no generation"): + slot.recycle() + with self.assertRaisesRegex(BusyError, "no active generation"): + slot.wait() + backend_slot.next_state = WorkState.PENDING + slot.start() + self.assertFalse(slot.test()) + with self.assertRaises(TransferTimeoutError): + slot.wait(0) + with self.assertRaisesRegex(BusyError, "non-terminal"): + slot.recycle() + self.assertTrue(slot.cancel()) + self.assertEqual(slot.state, WorkState.CANCELLED) + self.assertIsInstance(slot.error, TransferCancelledError) + with self.assertRaises(TransferCancelledError): + slot.wait() + slot.recycle() + + failure = RuntimeError("injected slot failure") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + slot.start() + self.assertEqual(slot.state, WorkState.FAILED) + self.assertIs(slot.error, failure) + with self.assertRaises(BackendFailureError) as context: + slot.wait() + self.assertIs(context.exception.__cause__, failure) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_and_wait_fuses_then_falls_back_or_times_out( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + self.assertTrue(source.capabilities.fused_request_slot_start_poll) + self.assertIs(slot.start_and_wait(), slot) + self.assertEqual(backend_slot.fused_calls, [(64, None)]) + self.assertEqual(backend_slot.poll_count, 0) + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + + self.assertIs(slot.start_and_wait(1 << 80), slot) + self.assertEqual(backend_slot.fused_calls[-1], (64, (1 << 63) - 1)) + slot.recycle() + + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.COMPLETED] + self.assertIs(slot.start_and_wait(1), slot) + max_polls, timeout_ns = backend_slot.fused_calls[-1] + self.assertEqual(max_polls, 64) + self.assertIsInstance(timeout_ns, int) + self.assertGreaterEqual(timeout_ns, 0) + self.assertLessEqual(timeout_ns, 1_000_000_000) + self.assertEqual(backend_slot.poll_count, 1) + slot.recycle() + + backend_slot.next_state = WorkState.RUNNING + with self.assertRaises(TransferTimeoutError): + slot.start_and_wait(0) + self.assertEqual(backend_slot.fused_calls[-1], (64, 0)) + self.assertTrue(slot.active) + self.assertEqual(backend_slot.poll_count, 1) + self.assertEqual(slot.state, WorkState.RUNNING) + self.assertTrue(slot.cancel()) + slot.recycle() + + generation = slot.generation + for timeout, error_type in ( + (True, TypeError), + ("one", TypeError), + (-1, ValueError), + (float("inf"), ValueError), + ): + with self.assertRaises(error_type): + slot.start_and_wait(timeout) # type: ignore[arg-type] + self.assertEqual(slot.generation, generation) + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_and_wait_uses_bounded_poll_chunks(self) -> None: + source = Endpoint( + "source", backend="test-request-slot-fused", thread_mode="single" + ) + destination = Endpoint( + "destination", backend="test-request-slot-fused", thread_mode="single" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [ + WorkState.RUNNING, + WorkState.RUNNING, + WorkState.COMPLETED, + ] + + with mock.patch.object( + source, + "_backend_request_slot_poll_bounded", + side_effect=AssertionError( + "SINGLE bounded polling must use direct cached dispatch" + ), + ): + self.assertIs(slot.start_and_wait(), slot) + self.assertEqual( + backend_slot.bounded_calls, + [ + ( + _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, + _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, + ) + ] + * 3, + ) + self.assertEqual(backend_slot.poll_count, 0) + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_poll_only_continuation_does_not_require_start_fusion( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-poll-only") + destination = Endpoint("destination", backend="test-request-slot-poll-only") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + + self.assertFalse(source.capabilities.fused_request_slot_start_poll) + self.assertTrue(source.capabilities.fused_request_slot_poll) + self.assertIs(slot.start_and_wait(), slot) + self.assertEqual(backend_slot.start_count, 1) + self.assertEqual(backend_slot.fused_calls, []) + self.assertEqual( + backend_slot.bounded_calls, + [ + ( + _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, + _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, + ) + ], + ) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_bounded_poll_timeout_leaves_generation_active( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.RUNNING] + + with mock.patch( + "torch.distributed._transfer._api.time.monotonic_ns", + side_effect=[0, 0, 0, 999_950_000, 1_000_000_000], + ): + with self.assertRaises(TransferTimeoutError): + slot.start_and_wait(1) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.RUNNING) + self.assertEqual( + backend_slot.bounded_calls, + [ + ( + _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + 50_000, + _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, + ) + ], + ) + self.assertTrue(slot.cancel()) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_bounded_poll_validates_provider_results(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + for result, message in ( + (None, "invalid transfer slot state"), + ("running", "invalid transfer slot state"), + (WorkState.PENDING, "regressed"), + ): + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [result] + with self.assertRaisesRegex(BackendFailureError, message): + slot.start_and_wait() + self.assertTrue(slot.active) + self.assertEqual(slot._state, WorkState.RUNNING) + backend_slot._state = WorkState.RUNNING + self.assertTrue(slot.cancel()) + slot.recycle() + + provider_error = RuntimeError("injected bounded-poll failure") + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_error = provider_error + with self.assertRaises(BackendFailureError) as context: + slot.start_and_wait() + self.assertIs(context.exception.__cause__, provider_error) + self.assertTrue(slot.active) + self.assertEqual(slot._state, WorkState.RUNNING) + backend_slot.bounded_error = None + backend_slot._state = WorkState.RUNNING + self.assertTrue(slot.cancel()) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_cancel_can_acquire_between_bounded_poll_chunks( + self, + ) -> None: + source = Endpoint( + "source", backend="test-request-slot-fused", thread_mode="multiple" + ) + destination = Endpoint( + "destination", backend="test-request-slot-fused", thread_mode="multiple" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.RUNNING] + between_chunks = threading.Event() + resume_waiter = threading.Event() + waiter_errors: list[BaseException] = [] + original_progress = Endpoint.progress + + def pause_between_chunks(self): + if self is source and not between_chunks.is_set(): + between_chunks.set() + if not resume_waiter.wait(2): + raise RuntimeError("bounded-poll cancellation test timed out") + return original_progress(self) + + def start_and_wait() -> None: + try: + slot.start_and_wait() + except BaseException as error: + waiter_errors.append(error) + + with mock.patch.object(Endpoint, "progress", pause_between_chunks): + waiter = threading.Thread(target=start_and_wait) + waiter.start() + self.assertTrue(between_chunks.wait(1)) + self.assertTrue(slot.cancel()) + resume_waiter.set() + waiter.join(2) + + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(waiter_errors), 1) + self.assertIsInstance(waiter_errors[0], TransferCancelledError) + self.assertEqual(len(backend_slot.bounded_calls), 1) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_bounded_wait_is_pinned_to_started_generation(self) -> None: + source = Endpoint( + "source", backend="test-request-slot-fused", thread_mode="multiple" + ) + destination = Endpoint( + "destination", backend="test-request-slot-fused", thread_mode="multiple" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + entered = threading.Event() + release = threading.Event() + errors: list[BaseException] = [] + original_wait_bounded = TransferSlot._wait_bounded_until + + def pause_before_bounded_wait(self, *args, **kwargs): + entered.set() + if not release.wait(2): + raise RuntimeError("bounded-poll generation test timed out") + return original_wait_bounded(self, *args, **kwargs) + + def run_waiter() -> None: + try: + slot.start_and_wait() + except BaseException as error: + errors.append(error) + + with mock.patch.object( + TransferSlot, "_wait_bounded_until", pause_before_bounded_wait + ): + waiter = threading.Thread(target=run_waiter) + waiter.start() + self.assertTrue(entered.wait(1)) + backend_slot._state = WorkState.COMPLETED + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + backend_slot.next_state = WorkState.COMPLETED + slot.start() + replacement_generation = slot.generation + release.set() + waiter.join(2) + + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIsInstance(errors[0], BusyError) + self.assertIn("changed while waiting", str(errors[0])) + self.assertEqual(slot.generation, replacement_generation) + self.assertEqual(backend_slot.bounded_calls, []) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_and_wait_validates_fused_provider_results( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + failure = RuntimeError("fused provider failure") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + with self.assertRaises(BackendFailureError) as context: + slot.start_and_wait() + self.assertIs(context.exception.__cause__, failure) + self.assertIs(slot.error, failure) + slot.recycle() + + backend_slot.next_state = WorkState.CANCELLED + backend_slot.next_error = None + with self.assertRaises(TransferCancelledError): + slot.start_and_wait() + self.assertIsInstance(slot.error, TransferCancelledError) + slot.recycle() + + for result, message in ( + (None, "without acknowledging"), + ("completed", "invalid fused request-slot state"), + ): + backend_slot.next_state = WorkState.COMPLETED + backend_slot.fused_result = result + with self.assertRaisesRegex(BackendFailureError, message): + slot.start_and_wait() + self.assertEqual(slot.state, WorkState.FAILED) + slot.recycle() + + backend_slot.fused_result = _RETURN_CURRENT_SLOT_STATE + backend_slot.fused_error = KeyboardInterrupt() + backend_slot.next_state = WorkState.COMPLETED + with self.assertRaises(BackendFailureError) as context: + slot.start_and_wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertEqual(slot.state, WorkState.FAILED) + slot.recycle() + + backend_slot.fused_error = None + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.PENDING] + with self.assertRaisesRegex(BackendFailureError, "regressed"): + slot.start_and_wait(1) + self.assertTrue(slot.active) + self.assertEqual(slot._state, WorkState.RUNNING) + backend_slot._state = WorkState.RUNNING + self.assertTrue(slot.cancel()) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_and_wait_fallback_ignores_fused_spi(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.start_and_poll = mock.Mock( + side_effect=AssertionError("fused SPI must not be entered") + ) + backend_slot.poll_bounded = mock.Mock( + side_effect=AssertionError("bounded-poll SPI must not be entered") + ) + + self.assertFalse(source.capabilities.fused_request_slot_start_poll) + self.assertFalse(source.capabilities.fused_request_slot_poll) + self.assertIs(slot.start_and_wait(), slot) + backend_slot.start_and_poll.assert_not_called() + backend_slot.poll_bounded.assert_not_called() + self.assertEqual(backend_slot.start_count, 1) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_source_compatible_fallback_recycles(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + self.assertFalse(source.capabilities.deferred_successful_request_slot_recycle) + self.assertEqual(slot.execute(), 1) + self.assertEqual(slot.generation, 1) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(backend_slot.start_count, 1) + self.assertEqual(backend_slot.recycle_count, 1) + with self.assertRaisesRegex(BusyError, "no active generation"): + _ = slot.completion_event + with self.assertRaisesRegex(BusyError, "no generation"): + slot.recycle() + + generation = slot.generation + for timeout, error_type in ( + (True, TypeError), + ("one", TypeError), + (-1, ValueError), + (float("inf"), ValueError), + ): + with self.assertRaises(error_type): + slot.execute(timeout) # type: ignore[arg-type] + self.assertEqual(slot.generation, generation) + self.assertFalse(slot.active) + + self.assertEqual(slot.execute(), 2) + self.assertEqual(backend_slot.recycle_count, 2) + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_deferred_recycle_preserves_explicit_api( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-deferred-recycle") + destination = Endpoint( + "destination", backend="test-request-slot-deferred-recycle" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + self.assertTrue(source.capabilities.deferred_successful_request_slot_recycle) + self.assertEqual(slot.execute(), 1) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertFalse(slot.test()) + self.assertIsNone(slot.error) + self.assertEqual(backend_slot.recycle_count, 0) + # The test provider models the retained, observable completion receipt. + self.assertEqual(backend_slot.state, WorkState.COMPLETED) + with self.assertRaisesRegex(BusyError, "no active generation"): + _ = slot.completion_event + + self.assertIs(slot.start_and_wait(), slot) + self.assertEqual(slot.generation, 2) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.COMPLETED) + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.start() + with self.assertRaisesRegex(BusyError, "recycle it before close"): + slot.close() + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + # Close after execute is legal even though the provider still owns its + # deferred successful receipt. + self.assertEqual(slot.execute(), 3) + self.assertEqual(backend_slot.recycle_count, 1) + slot.close() + self.assertTrue(slot.closed) + self.assertEqual(backend_slot.close_count, 1) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_deferred_fused_immediate_success(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + with ( + mock.patch.object( + source, + "_backend_request_slot_start_and_poll", + side_effect=AssertionError("SINGLE execute must use cached dispatch"), + ), + mock.patch.object( + _api, + "_transfer_slot_deadline_ns", + side_effect=AssertionError("timeout=None must bypass validation"), + ), + mock.patch.object( + TransferSlot, + "_check_access", + side_effect=AssertionError("fast execute must inline access checks"), + ), + mock.patch.object( + TransferSlot, + "_begin_generation_unlocked", + side_effect=AssertionError("fast execute must inline generation begin"), + ), + mock.patch.object( + TransferSlot, + "_start_and_poll_unlocked", + side_effect=AssertionError("fast execute must inline fused start"), + ), + mock.patch.object( + TransferSlot, + "_observe_backend_state_unlocked", + side_effect=AssertionError("fast execute must inline completion"), + ), + ): + self.assertEqual(slot.execute(), 1) + self.assertEqual(slot.execute(), 2) + self.assertEqual(backend_slot.fused_calls, [(64, None), (64, None)]) + self.assertEqual(backend_slot.poll_count, 0) + self.assertEqual(backend_slot.recycle_count, 0) + self.assertEqual(slot.generation, 2) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_exact_running_uses_direct_continuation( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + observe = TransferSlot._observe_backend_state_unlocked + + def reject_initial_running(self, state): + if state is WorkState.RUNNING: + raise AssertionError("initial RUNNING must use direct continuation") + return observe(self, state) + + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + with mock.patch.object( + TransferSlot, + "_observe_backend_state_unlocked", + reject_initial_running, + ): + self.assertEqual(slot.execute(), 1) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(len(backend_slot.bounded_calls), 1) + + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.PENDING] + with ( + mock.patch.object( + TransferSlot, + "_observe_backend_state_unlocked", + reject_initial_running, + ), + self.assertRaisesRegex(BackendFailureError, "regressed"), + ): + slot.execute() + self.assertTrue(slot.active) + self.assertTrue(slot._backend_generation_active) + self.assertIs(slot._state, WorkState.RUNNING) + backend_slot._state = WorkState.RUNNING + self.assertTrue(slot.cancel()) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_test_recycles_only_success(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + with mock.patch.object( + TransferSlot, + "_start_and_poll_unlocked", + side_effect=AssertionError("exact start/test must use cached dispatch"), + ): + self.assertTrue(slot.start_and_test()) + self.assertFalse(slot.active) + self.assertEqual(slot.generation, 1) + self.assertEqual(backend_slot.fused_calls, [(0, None)]) + self.assertEqual(backend_slot.recycle_count, 0) + + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.RUNNING, WorkState.COMPLETED] + with mock.patch.object( + TransferSlot, + "_observe_backend_state_unlocked", + side_effect=AssertionError( + "exact RUNNING start/test must not reclassify provider state" + ), + ): + self.assertFalse(slot.start_and_test()) + self.assertTrue(slot.active) + with ( + mock.patch.object( + TransferSlot, + "_check_access", + side_effect=AssertionError("exact continuation must inline access"), + ), + mock.patch.object( + TransferSlot, + "_refresh_unlocked", + side_effect=AssertionError("exact continuation must poll directly"), + ), + mock.patch.object( + TransferSlot, + "_finish_execute_success_unlocked", + side_effect=AssertionError("exact continuation must finish inline"), + ), + mock.patch.object( + TransferSlot, + "_observe_backend_state_unlocked", + side_effect=AssertionError( + "exact continuation must inline common provider states" + ), + ), + ): + self.assertFalse(slot.test_and_recycle()) + self.assertTrue(slot.test_and_recycle()) + self.assertFalse(slot.active) + self.assertEqual(backend_slot.recycle_count, 0) + with self.assertRaisesRegex(BusyError, "no active generation"): + slot.test_and_recycle() + + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = RuntimeError("injected failure") + with self.assertRaises(BackendFailureError): + slot.start_and_test() + self.assertTrue(slot.active) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_cached_success_recycles_without_another_poll(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.COMPLETED] + self.assertFalse(slot.start_and_test()) + self.assertTrue(slot.test()) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.COMPLETED) + polls_before_recycle = backend_slot.poll_count + with mock.patch.object( + slot, + "_backend_poll_state", + side_effect=AssertionError("cached completion must not poll again"), + ): + self.assertTrue(slot.test_and_recycle()) + self.assertEqual(backend_slot.poll_count, polls_before_recycle) + self.assertEqual(backend_slot.recycle_count, 0) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_serialized_fallback_start_test_is_one_transaction( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-blocking-observation", + thread_mode="serialized", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-blocking-observation", + thread_mode="serialized", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.block_next_observation = True + results: dict[str, bool] = {} + errors: list[BaseException] = [] + second_entered = threading.Event() + second_done = threading.Event() + + def start_and_test(label: str) -> None: + if label == "second": + second_entered.set() + try: + results[label] = slot.start_and_test() + except BaseException as error: + errors.append(error) + finally: + if label == "second": + second_done.set() + + first = threading.Thread(target=start_and_test, args=("first",)) + second = threading.Thread(target=start_and_test, args=("second",)) + first.start() + self.assertTrue(backend_slot.observation_entered.wait(1)) + second.start() + self.assertTrue(second_entered.wait(1)) + try: + self.assertFalse(second_done.wait(0.05)) + self.assertEqual(backend_slot.start_count, 1) + self.assertEqual(slot.generation, 1) + self.assertTrue(slot.active) + finally: + backend_slot.observation_release.set() + first.join(2) + second.join(2) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(results, {"first": True, "second": True}) + self.assertEqual(backend_slot.start_count, 2) + self.assertEqual(backend_slot.recycle_count, 2) + self.assertEqual(slot.generation, 2) + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_test_retains_cancel_and_failure(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + backend_slot.next_state = WorkState.RUNNING + self.assertFalse(slot.start_and_test()) + self.assertTrue(slot.cancel()) + with self.assertRaises(TransferCancelledError): + slot.test_and_recycle() + self.assertTrue(slot.active) + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.start_and_test() + with self.assertRaisesRegex(BusyError, "recycle it before close"): + slot.close() + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + failure = RuntimeError("retained start/test failure") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + with self.assertRaises(BackendFailureError) as context: + slot.start_and_test() + self.assertIs(context.exception.__cause__, failure) + self.assertTrue(slot.active) + with self.assertRaises(BackendFailureError) as context: + slot.test_and_recycle() + self.assertIs(context.exception.__cause__, failure) + self.assertTrue(slot.active) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 2) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_test_deferred_success_closes_without_recycle( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + self.assertTrue(slot.start_and_test()) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(backend_slot.recycle_count, 0) + slot.close() + self.assertTrue(slot.closed) + self.assertEqual(backend_slot.recycle_count, 0) + self.assertEqual(backend_slot.close_count, 1) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_test_interruption_boundaries_are_recoverable( + self, + ) -> None: + class InterruptingIdleCommitTransferSlot(TransferSlot): + __slots__ = ("interrupt_idle_commit",) + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "TransferSlot", InterruptingIdleCommitTransferSlot + ): + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + backend_slot.fused_error = KeyboardInterrupt() + with self.assertRaises(BackendFailureError) as context: + slot.start_and_test() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.FAILED) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + backend_slot.fused_error = None + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.start_and_test() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot.generation, 2) + self.assertEqual(backend_slot.recycle_count, 1) + self.assertTrue(slot.start_and_test()) + self.assertEqual(slot.generation, 3) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_slot_start_test_snapshots_notification(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-notification", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-notification", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot(request_local_notifications=True) + backend_slot = source._backend.request_slots[0] + + payload = bytearray(b"A\x00\xff") + with mock.patch.object( + RequestLocalNotificationSlot, + "_start_and_poll_with_notification_unlocked", + side_effect=AssertionError( + "exact notification start/test must use cached dispatch" + ), + ): + self.assertTrue(slot.start_and_test(notification=payload)) + payload[:] = b"BAD" + self.assertEqual( + backend_slot.notification_fused_calls, + [(b"A\x00\xff", 0, None)], + ) + self.assertTrue(slot.start_and_test(notification=None)) + self.assertEqual( + backend_slot.notification_fused_calls[-1], + (None, 0, None), + ) + self.assertFalse(slot.active) + self.assertEqual(backend_slot.recycle_count, 0) + + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.COMPLETED] + with mock.patch.object( + RequestLocalNotificationSlot, + "_observe_backend_state_unlocked", + side_effect=AssertionError( + "exact notification RUNNING start/test must not reclassify state" + ), + ): + self.assertFalse(slot.start_and_test(notification=b"pending")) + self.assertTrue(slot.active) + self.assertTrue(slot.test_and_recycle()) + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_slot_execute_exact_running_uses_direct_continuation( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-notification", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-notification", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot(request_local_notifications=True) + backend_slot = source._backend.request_slots[0] + observe = RequestLocalNotificationSlot._observe_backend_state_unlocked + + def reject_initial_running(self, state): + if state is WorkState.RUNNING: + raise AssertionError("initial RUNNING must use direct continuation") + return observe(self, state) + + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + with mock.patch.object( + RequestLocalNotificationSlot, + "_observe_backend_state_unlocked", + reject_initial_running, + ): + self.assertEqual(slot.execute(notification=b"ready"), 1) + self.assertEqual( + backend_slot.notification_fused_calls, + [(b"ready", _api._FUSED_REQUEST_SLOT_MAX_POLLS, None)], + ) + self.assertEqual(len(backend_slot.bounded_calls), 1) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertIs( + slot._notification_fast_execute, + slot._backend_start_and_poll_with_notification, + ) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_slot_interrupted_completion_cannot_reuse_stale_state( + self, + ) -> None: + class InterruptingNotificationSlot(RequestLocalNotificationSlot): + __slots__ = ("interrupt_idle_commit",) + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint( + "source", backend="test-request-slot-notification", thread_mode="single" + ) + destination = Endpoint( + "destination", + backend="test-request-slot-notification", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "RequestLocalNotificationSlot", InterruptingNotificationSlot + ): + slot = bound.create_slot(request_local_notifications=True) + backend_slot = source._backend.request_slots[0] + notification_fast_execute = slot._notification_fast_execute + self.assertIsNotNone(notification_fast_execute) + + backend_slot.next_state = WorkState.RUNNING + self.assertFalse(slot.start_and_test(notification=b"first")) + backend_slot.poll_states = [WorkState.COMPLETED] + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.test_and_recycle() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertIsNone(slot._state) + self.assertIsNone(slot._start_error) + self.assertIs( + slot._notification_fast_execute, + notification_fast_execute, + ) + + prior_poll_count = backend_slot.poll_count + backend_slot.next_state = WorkState.RUNNING + backend_slot.fused_error = KeyboardInterrupt() + backend_slot.poll_states = [WorkState.COMPLETED] + with self.assertRaises(BackendFailureError) as context: + slot.start_and_test(notification=b"second") + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertEqual(backend_slot.poll_count, prior_poll_count + 1) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.FAILED) + + slot.recycle() + self.assertFalse(slot.active) + self.assertIs( + slot._notification_fast_execute, + slot._backend_start_and_poll_with_notification, + ) + backend_slot.fused_error = None + backend_slot.next_state = WorkState.COMPLETED + self.assertTrue(slot.start_and_test(notification=None)) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_request_local_slot_interrupted_bounded_execute_uses_safe_fallback( + self, + ) -> None: + class InterruptingNotificationSlot(RequestLocalNotificationSlot): + __slots__ = ("interrupt_idle_commit",) + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint( + "source", backend="test-request-slot-notification", thread_mode="single" + ) + destination = Endpoint( + "destination", + backend="test-request-slot-notification", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "RequestLocalNotificationSlot", InterruptingNotificationSlot + ): + slot = bound.create_slot(request_local_notifications=True) + backend_slot = source._backend.request_slots[0] + + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.execute(notification=b"bounded") + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot._state, WorkState.COMPLETED) + self.assertIsNone(slot._notification_fast_execute) + + backend_slot.next_state = WorkState.COMPLETED + self.assertEqual(slot.execute(notification=b"fallback"), 2) + self.assertFalse(slot.active) + self.assertIsNone(slot._state) + self.assertIs( + slot._notification_fast_execute, + slot._backend_start_and_poll_with_notification, + ) + + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = RuntimeError("notification failure") + with self.assertRaises(BackendFailureError): + slot.start_and_test(notification=b"failure") + self.assertIsNotNone(slot._notification_fast_execute) + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.recycle() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot._state, WorkState.FAILED) + self.assertIsNone(slot._notification_fast_execute) + + backend_slot.next_state = WorkState.COMPLETED + backend_slot.next_error = None + self.assertTrue(slot.start_and_test(notification=None)) + self.assertIs( + slot._notification_fast_execute, + slot._backend_start_and_poll_with_notification, + ) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_single_fused_fast_path_preserves_guards( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + wrong_thread_errors: list[BaseException] = [] + + def execute_from_wrong_thread() -> None: + try: + slot.execute() + except BaseException as error: + wrong_thread_errors.append(error) + + wrong_thread = threading.Thread(target=execute_from_wrong_thread) + wrong_thread.start() + wrong_thread.join(1) + self.assertFalse(wrong_thread.is_alive()) + self.assertEqual(len(wrong_thread_errors), 1) + self.assertIsInstance(wrong_thread_errors[0], UnsupportedError) + self.assertEqual(slot.generation, 0) + self.assertFalse(slot.active) + self.assertEqual(backend_slot.fused_calls, []) + + backend_slot.next_state = WorkState.COMPLETED + backend_slot.fused_error = KeyboardInterrupt() + with self.assertRaises(BackendFailureError) as context: + slot.execute() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertEqual(slot.generation, 1) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.FAILED) + slot.recycle() + + backend_slot.fused_error = None + slot.start() + fused_call_count = len(backend_slot.fused_calls) + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.execute() + self.assertEqual(len(backend_slot.fused_calls), fused_call_count) + self.assertEqual(slot.generation, 2) + slot.recycle() + + self.assertEqual(slot.execute(), 3) + slot.close() + with self.assertRaises(ClosedError): + slot.execute() + self.assertEqual(slot.generation, 3) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_bounded_success_is_atomic_in_multiple( + self, + ) -> None: + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="multiple", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="multiple", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + finish_entered = threading.Event() + finish_release = threading.Event() + observer_done = threading.Event() + execute_results: list[int] = [] + errors: list[BaseException] = [] + observed_states: list[WorkState | None] = [] + original_finish = TransferSlot._finish_execute_success_unlocked + + def pause_in_finish(self, expected_generation): + finish_entered.set() + if not finish_release.wait(2): + raise RuntimeError("execute atomicity test timed out") + return original_finish(self, expected_generation) + + def run_execute() -> None: + try: + execute_results.append(slot.execute()) + except BaseException as error: + errors.append(error) + + def observe_state() -> None: + try: + observed_states.append(slot.state) + except BaseException as error: + errors.append(error) + finally: + observer_done.set() + + with mock.patch.object( + TransferSlot, "_finish_execute_success_unlocked", pause_in_finish + ): + executor = threading.Thread(target=run_execute) + executor.start() + self.assertTrue(finish_entered.wait(1)) + observer = threading.Thread(target=observe_state) + observer.start() + self.assertFalse(observer_done.wait(0.05)) + finish_release.set() + executor.join(2) + observer.join(2) + + self.assertFalse(executor.is_alive()) + self.assertFalse(observer.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(execute_results, [1]) + self.assertEqual(observed_states, [None]) + self.assertEqual(len(backend_slot.bounded_calls), 1) + self.assertEqual(backend_slot.recycle_count, 0) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_retains_unsuccessful_generations(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused-deferred-recycle") + destination = Endpoint( + "destination", backend="test-request-slot-fused-deferred-recycle" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + failure = RuntimeError("execute provider failure") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + with self.assertRaises(BackendFailureError) as context: + slot.execute() + self.assertIs(context.exception.__cause__, failure) + self.assertTrue(slot.active) + self.assertEqual(slot.generation, 1) + slot.recycle() + + backend_slot.next_state = WorkState.CANCELLED + backend_slot.next_error = None + with self.assertRaises(TransferCancelledError): + slot.execute() + self.assertTrue(slot.active) + self.assertEqual(slot.generation, 2) + slot.recycle() + + backend_slot.next_state = WorkState.RUNNING + with self.assertRaises(TransferTimeoutError): + slot.execute(0) + self.assertTrue(slot.active) + self.assertEqual(slot.generation, 3) + self.assertEqual(slot.state, WorkState.RUNNING) + self.assertTrue(slot.cancel()) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 3) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_execute_interruption_preserves_a_safe_boundary( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-fused-deferred-recycle") + destination = Endpoint( + "destination", backend="test-request-slot-fused-deferred-recycle" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + original_finish = TransferSlot._finish_execute_success_unlocked + + def interrupt_before_finish(self, expected_generation): + raise KeyboardInterrupt + + with mock.patch.object( + TransferSlot, + "_finish_execute_success_unlocked", + interrupt_before_finish, + ): + with self.assertRaises(KeyboardInterrupt): + slot.execute() + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + def interrupt_after_finish(self, expected_generation): + original_finish(self, expected_generation) + raise KeyboardInterrupt + + with mock.patch.object( + TransferSlot, + "_finish_execute_success_unlocked", + interrupt_after_finish, + ): + with self.assertRaises(KeyboardInterrupt): + slot.execute() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot.generation, 2) + self.assertEqual(backend_slot.recycle_count, 1) + + self.assertEqual(slot.execute(), 3) + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_deferred_execute_idle_boundaries_are_recoverable( + self, + ) -> None: + class InterruptingIdleCommitTransferSlot(TransferSlot): + __slots__ = ("interrupt_fast_restore", "interrupt_idle_commit") + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_fast_restore = False + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + if ( + name == "_fast_execute" + and value is not None + and getattr(self, "interrupt_fast_restore", False) + ): + object.__setattr__(self, "interrupt_fast_restore", False) + raise KeyboardInterrupt + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "TransferSlot", InterruptingIdleCommitTransferSlot + ): + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + fast_execute = slot._fast_execute + self.assertIs(fast_execute, slot._backend_start_and_poll) + self.assertTrue(slot._caller_ready) + + def interrupt_before_provider_entry(**_kwargs) -> WorkState: + raise KeyboardInterrupt + + slot._fast_execute = interrupt_before_provider_entry + with self.assertRaises(BackendFailureError) as context: + slot.execute() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertTrue(slot.active) + self.assertEqual(slot.state, WorkState.FAILED) + self.assertEqual(backend_slot.fused_calls, []) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 0) + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.execute() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertIsNone(slot.error) + self.assertIsNone(slot._state) + self.assertEqual(slot.generation, 2) + self.assertEqual(backend_slot.state, WorkState.COMPLETED) + self.assertIs(slot._fast_execute, fast_execute) + self.assertEqual(slot.execute(), 3) + self.assertEqual(slot.execute(), 4) + self.assertIsNone(slot._state) + self.assertEqual(backend_slot.recycle_count, 0) + + # A timed call uses the generic success finalizer. It disables exact + # execution before public idle, so an interrupt at that commit leaves a + # dirty but fast-ineligible slot that the next generic call repairs. + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.execute(1.0) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(slot._state, WorkState.COMPLETED) + self.assertIsNone(slot._fast_execute) + + # Publishing exact eligibility is the last private-idle store. An + # interrupt just before it leaves a completely clean generic fallback; + # one successful fallback restores the exact path for its successor. + slot.interrupt_fast_restore = True + with self.assertRaises(KeyboardInterrupt): + slot.execute(1.0) + self.assertFalse(slot.active) + self.assertIsNone(slot._state) + self.assertIsNone(slot._error) + self.assertIsNone(slot._start_error) + self.assertFalse(slot._backend_generation_active) + self.assertIsNone(slot._fast_execute) + self.assertEqual(slot.execute(), 7) + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + self.assertEqual(slot.execute(), 8) + + slot.close() + self.assertIsNone(slot._fast_execute) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_dirty_idle_repairs_ambiguous_failure_before_begin( + self, + ) -> None: + class InterruptingIdleCommitTransferSlot(TransferSlot): + __slots__ = ("interrupt_fast_restore", "interrupt_idle_commit") + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_fast_restore = False + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + if ( + name == "_fast_execute" + and value is not None + and getattr(self, "interrupt_fast_restore", False) + ): + object.__setattr__(self, "interrupt_fast_restore", False) + raise KeyboardInterrupt + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint( + "source", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-fused-deferred-recycle", + thread_mode="single", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "TransferSlot", InterruptingIdleCommitTransferSlot + ): + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + backend_slot.next_state = WorkState.COMPLETED + backend_slot.fused_error = KeyboardInterrupt() + with self.assertRaises(BackendFailureError): + slot.execute() + self.assertEqual(slot.state, WorkState.FAILED) + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.recycle() + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertIsNone(slot.error) + self.assertEqual(slot._state, WorkState.FAILED) + self.assertIsNotNone(slot._start_error) + self.assertIsNone(slot._fast_execute) + + backend_slot.fused_error = None + backend_slot.next_state = WorkState.RUNNING + backend_slot.bounded_results = [WorkState.COMPLETED] + self.assertEqual(slot.execute(), 2) + self.assertFalse(slot.active) + self.assertIsNone(slot.error) + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + + # CANCELLED is never folded. Its explicit recycle follows the same + # disable-clean-restore protocol; interruption before the last restore + # leaves a clean generic fallback and cannot expose stale diagnostics. + backend_slot.next_state = WorkState.CANCELLED + with self.assertRaises(TransferCancelledError): + slot.execute() + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + slot.interrupt_fast_restore = True + with self.assertRaises(KeyboardInterrupt): + slot.recycle() + self.assertFalse(slot.active) + self.assertIsNone(slot._state) + self.assertIsNone(slot._error) + self.assertIsNone(slot._start_error) + self.assertFalse(slot._backend_generation_active) + self.assertIsNone(slot._fast_execute) + + backend_slot.next_state = WorkState.COMPLETED + self.assertEqual(slot.execute(), 4) + self.assertIs(slot._fast_execute, slot._backend_start_and_poll) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_fused_fallback_releases_slot_lock(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-start-fused-only", + thread_mode="multiple", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-start-fused-only", + thread_mode="multiple", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.start_entered = threading.Event() + errors: list[BaseException] = [] + + def start_and_wait() -> None: + try: + slot.start_and_wait(2) + except BaseException as error: + errors.append(error) + + waiter = threading.Thread(target=start_and_wait) + waiter.start() + self.assertTrue(backend_slot.start_entered.wait(1)) + self.assertTrue(slot.cancel()) + waiter.join(2) + + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIsInstance(errors[0], TransferCancelledError) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_waits_are_pinned_to_the_started_generation(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-start-fused-only", + thread_mode="multiple", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-start-fused-only", + thread_mode="multiple", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + original_wait_until = TransferSlot._wait_until + + def replace_generation_while_waiter_is_paused(start_and_wait: bool) -> None: + entered = threading.Event() + release = threading.Event() + errors: list[BaseException] = [] + + def pause_before_first_observation(self, *args, **kwargs): + entered.set() + if not release.wait(2): + raise RuntimeError("generation-pinning test timed out") + return original_wait_until(self, *args, **kwargs) + + if start_and_wait: + backend_slot.next_state = WorkState.RUNNING + operation = slot.start_and_wait + else: + backend_slot.next_state = WorkState.COMPLETED + slot.start() + operation = slot.wait + + def run_waiter() -> None: + try: + operation() + except BaseException as error: + errors.append(error) + + with mock.patch.object( + TransferSlot, "_wait_until", pause_before_first_observation + ): + waiter = threading.Thread(target=run_waiter) + waiter.start() + self.assertTrue(entered.wait(1)) + backend_slot._state = WorkState.COMPLETED + self.assertEqual(slot.state, WorkState.COMPLETED) + slot.recycle() + backend_slot.next_state = WorkState.COMPLETED + slot.start() + replacement_generation = slot.generation + release.set() + waiter.join(2) + + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIsInstance(errors[0], BusyError) + self.assertIn("changed while waiting", str(errors[0])) + self.assertEqual(slot.generation, replacement_generation) + slot.recycle() + + replace_generation_while_waiter_is_paused(start_and_wait=False) + replace_generation_while_waiter_is_paused(start_and_wait=True) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_start_and_wait_fallback_remains_spin_oriented(self) -> None: + source = Endpoint("source", backend="test-request-slot-start-fused-only") + destination = Endpoint( + "destination", backend="test-request-slot-start-fused-only" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + backend_slot.poll_states = [WorkState.RUNNING] * 130 + [WorkState.COMPLETED] + + with mock.patch("torch.distributed._transfer._api.time.sleep") as sleep: + self.assertIs(slot.start_and_wait(), slot) + sleep.assert_not_called() + self.assertEqual(backend_slot.poll_count, 131) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_wait_retains_cooperative_backoff(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.next_state = WorkState.RUNNING + slot.start() + + def complete_on_yield(delay: float) -> None: + self.assertEqual(delay, 0) + backend_slot._state = WorkState.COMPLETED + + with mock.patch( + "torch.distributed._transfer._api.time.sleep", + side_effect=complete_on_yield, + ) as sleep: + self.assertIs(slot.wait(), slot) + sleep.assert_called_once_with(0) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_rejects_concurrent_start(self) -> None: + source = Endpoint("source", backend="test-request-slot", thread_mode="multiple") + destination = Endpoint( + "destination", backend="test-request-slot", thread_mode="multiple" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.start_entered = threading.Event() + backend_slot.start_release = threading.Event() + errors: list[BaseException] = [] + + def start_slot() -> None: + try: + slot.start() + except BaseException as error: + errors.append(error) + + first = threading.Thread(target=start_slot) + second = threading.Thread(target=start_slot) + first.start() + self.assertTrue(backend_slot.start_entered.wait(1)) + second.start() + backend_slot.start_release.set() + first.join(2) + second.join(2) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIsInstance(errors[0], BusyError) + self.assertEqual(backend_slot.start_count, 1) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_close_is_linearized_with_fused_start(self) -> None: + source = Endpoint( + "source", backend="test-request-slot-fused", thread_mode="multiple" + ) + destination = Endpoint( + "destination", backend="test-request-slot-fused", thread_mode="multiple" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + backend_slot.start_entered = threading.Event() + backend_slot.start_release = threading.Event() + start_errors: list[BaseException] = [] + close_errors: list[BaseException] = [] + + def start_slot() -> None: + try: + slot.start_and_wait() + except BaseException as error: + start_errors.append(error) + + def close_slot() -> None: + try: + slot.close() + except BaseException as error: + close_errors.append(error) + + starter = threading.Thread(target=start_slot) + closer = threading.Thread(target=close_slot) + starter.start() + self.assertTrue(backend_slot.start_entered.wait(1)) + closer.start() + backend_slot.start_release.set() + starter.join(2) + closer.join(2) + + self.assertFalse(starter.is_alive()) + self.assertFalse(closer.is_alive()) + self.assertEqual(start_errors, []) + self.assertEqual(len(close_errors), 1) + self.assertIsInstance(close_errors[0], BusyError) + self.assertTrue(slot.active) + self.assertFalse(slot.closed) + + slot.recycle() + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_interrupted_before_provider_ack_is_recyclable(self) -> None: + source = Endpoint("source", backend="test-request-slot-unacknowledged-start") + destination = Endpoint( + "destination", backend="test-request-slot-unacknowledged-start" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + self.assertIsNone(slot.start()) + self.assertEqual(slot.state, WorkState.FAILED) + self.assertIsInstance(slot.error, KeyboardInterrupt) + with self.assertRaises(BackendFailureError) as context: + slot.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + slot.recycle() + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_normal_return_without_provider_ack_is_recyclable( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-no-ack") + destination = Endpoint("destination", backend="test-request-slot-no-ack") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + self.assertIsNone(slot.start()) + self.assertEqual(slot.state, WorkState.FAILED) + with self.assertRaisesRegex(BackendFailureError, "without acknowledging"): + slot.wait() + slot.recycle() + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_fused_interruption_before_ack_is_recyclable(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused-unacknowledged") + destination = Endpoint( + "destination", backend="test-request-slot-fused-unacknowledged" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + with self.assertRaises(BackendFailureError) as context: + slot.start_and_wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + self.assertEqual(slot.state, WorkState.FAILED) + self.assertTrue(slot.active) + slot.recycle() + self.assertFalse(slot.active) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_single_mode_rejects_every_wrong_thread_operation( + self, + ) -> None: + source = Endpoint( + "source", backend="test-request-slot-fused", thread_mode="single" + ) + destination = Endpoint( + "destination", backend="test-request-slot-fused", thread_mode="single" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + def wrong_thread_error(function) -> BaseException: + errors: list[BaseException] = [] + + def invoke() -> None: + try: + function() + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=invoke) + thread.start() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(len(errors), 1) + return errors[0] + + for function in ( + slot.start, + slot.start_and_wait, + lambda: slot.state, + slot.test, + slot.recycle, + slot.close, + ): + self.assertIsInstance(wrong_thread_error(function), UnsupportedError) + + slot.start() + slot.recycle() + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_lifetime_blockers_and_recovery_snapshot(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + self.assertEqual(source.open_transfer_slots, (slot,)) + with self.assertRaisesRegex(BusyError, "request slots"): + bound.close() + with self.assertRaisesRegex(BusyError, "transfer slot"): + source.close() + slot.start() + with self.assertRaisesRegex(BusyError, "unrecycled generation"): + slot.close() + slot.recycle() + slot.close() + self.assertEqual(source.open_transfer_slots, ()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_host_wait_runs_on_every_fused_start(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local.region() + local_region._device = "cuda:0" + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + actions: list[str] = [] + synchronize_bind_count = 0 + + class SynchronizeDescriptor: + def __get__(self, instance: object, owner: object) -> object: + nonlocal synchronize_bind_count + synchronize_bind_count += 1 + return instance._synchronize + + class FakeEvent: + device = torch.device("cuda:0") + synchronize = SynchronizeDescriptor() + + def _synchronize(self) -> None: + actions.append("wait") + + with ( + mock.patch.object(torch.cuda, "Event", FakeEvent), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=False + ) as capture_query, + ): + bound = plan.bind( + TransferOp.WRITE, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(), + ) + slot = bound.create_slot() + self.assertEqual(synchronize_bind_count, 1) + for _ in range(2): + slot.start_and_wait() + slot.recycle() + + self.assertEqual(actions, ["wait", "wait"]) + self.assertEqual(synchronize_bind_count, 1) + capture_query.assert_not_called() + slot.close() + self.assertIsNone(slot._ready_event_synchronize) + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_cached_ready_event_failure_precedes_generation(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local.region() + local_region._device = "cuda:0" + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + + class FailingEvent: + device = torch.device("cuda:0") + + def synchronize(self) -> None: + raise RuntimeError("injected CUDA wait failure") + + with mock.patch.object(torch.cuda, "Event", FailingEvent): + bound = plan.bind( + TransferOp.WRITE, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FailingEvent(), + ) + slot = bound.create_slot() + with self.assertRaisesRegex(BackendFailureError, "CUDA readiness"): + slot.start_and_wait() + + self.assertEqual(slot.generation, 0) + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertEqual(source._backend.request_slots[0].start_count, 0) + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_host_wait_stream_reuses_one_event(self) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local.region() + local_region._device = "cuda:0" + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + actions: list[tuple[str, int]] = [] + event_count = 0 + + class FakeEvent: + device = torch.device("cuda:0") + + def __init__(self) -> None: + nonlocal event_count + event_count += 1 + self.event_id = event_count + + def record(self, stream: object) -> None: + self.assert_stream(stream) + actions.append(("record", self.event_id)) + + def assert_stream(self, stream: object) -> None: + if not isinstance(stream, FakeStream): + raise AssertionError("readiness event recorded on wrong stream") + + def synchronize(self) -> None: + actions.append(("wait", self.event_id)) + + class FakeStream: + device = torch.device("cuda:0") + + def is_capturing(self) -> bool: + actions.append(("capture", event_count)) + return False + + with ( + mock.patch.object(torch, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Event", FakeEvent), + ): + bound = plan.bind( + TransferOp.WRITE, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + stream=FakeStream(), + ) + slot = bound.create_slot() + self.assertEqual(event_count, 1) + self.assertEqual(slot.execute(), 1) + self.assertEqual(slot.execute(), 2) + + self.assertEqual( + actions, + [ + ("capture", 1), + ("record", 1), + ("wait", 1), + ("capture", 1), + ("record", 1), + ("wait", 1), + ], + ) + self.assertEqual(source._backend.request_slots[0].start_count, 2) + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_reusable_stream_event_failure_precedes_generation( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-fused") + destination = Endpoint("destination", backend="test-request-slot-fused") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local.region() + local_region._device = "cuda:0" + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + event_count = 0 + wait_count = 0 + + class FakeEvent: + device = torch.device("cuda:0") + + def __init__(self) -> None: + nonlocal event_count + event_count += 1 + + def record(self, _stream: object) -> None: + return None + + def synchronize(self) -> None: + nonlocal wait_count + wait_count += 1 + if wait_count == 1: + raise RuntimeError("injected reusable-event wait failure") + + class FakeStream: + device = torch.device("cuda:0") + + def is_capturing(self) -> bool: + return False + + with ( + mock.patch.object(torch, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Event", FakeEvent), + ): + bound = plan.bind( + TransferOp.WRITE, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + stream=FakeStream(), + ) + slot = bound.create_slot() + with self.assertRaisesRegex(BackendFailureError, "CUDA readiness"): + slot.execute() + self.assertEqual(slot.generation, 0) + self.assertFalse(slot.active) + self.assertEqual(source._backend.request_slots[0].start_count, 0) + self.assertEqual(slot.execute(), 1) + + self.assertEqual(event_count, 1) + self.assertEqual(wait_count, 2) + self.assertEqual(source._backend.request_slots[0].start_count, 1) + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_creation_interruption_rolls_back_or_is_recoverable( + self, + ) -> None: + for backend_name in ( + "test-request-slot-before-adoption-interrupt", + "test-request-slot-create-interrupt", + ): + source = Endpoint("source", backend=backend_name) + destination = Endpoint("destination", backend=backend_name) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region()], remote=[peer.region("buffer")] + ) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaises(KeyboardInterrupt): + bound.create_slot() + self.assertEqual(source.open_transfer_slots, ()) + self.assertFalse(bound._slots) + if backend_name == "test-request-slot-create-interrupt": + self.assertEqual(source._backend.created_slot.close_count, 1) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + source = Endpoint("source", backend="test-request-slot-cleanup-interrupt") + destination = Endpoint( + "destination", backend="test-request-slot-cleanup-interrupt" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + + with self.assertRaisesRegex(BackendFailureError, "cleanup is pending"): + bound.create_slot() + (slot,) = source.open_transfer_slots + with self.assertRaisesRegex(BackendFailureError, "cleanup only"): + slot.start() + slot.close() + self.assertEqual(source.open_transfer_slots, ()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_creation_publication_interrupt_has_no_hidden_blocker( + self, + ) -> None: + class InterruptAfterAdd(set): + def add(self, value) -> None: + super().add(value) + raise KeyboardInterrupt + + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + bound._slots = InterruptAfterAdd(bound._slots) + + with self.assertRaises(KeyboardInterrupt): + bound.create_slot() + self.assertEqual(source.open_transfer_slots, ()) + self.assertFalse(bound._slots) + self.assertEqual(source._backend.request_slots, []) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_close_commit_interruptions_are_retryable(self) -> None: + class InterruptingSet(set): + def __init__(self, values, *, before: bool) -> None: + super().__init__(values) + self.before = before + self.armed = True + + def discard(self, value) -> None: + if self.armed: + self.armed = False + if self.before: + raise KeyboardInterrupt + super().discard(value) + raise KeyboardInterrupt + super().discard(value) + + for boundary in ("endpoint-after", "bound-before", "bound-after"): + with self.subTest(boundary=boundary): + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region()], remote=[peer.region("buffer")] + ) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + if boundary == "endpoint-after": + source._transfer_slots = InterruptingSet( + source._transfer_slots, before=False + ) + else: + bound._slots = InterruptingSet( + bound._slots, before=boundary == "bound-before" + ) + + with self.assertRaises(KeyboardInterrupt): + slot.close() + if boundary == "endpoint-after": + self.assertTrue(slot.closed) + self.assertNotIn(slot, bound._slots) + else: + self.assertTrue(slot.closed) + slot.close() + self.assertTrue(slot.closed) + self.assertIsNone(slot._bound) + self.assertNotIn(slot, source._transfer_slots) + self.assertNotIn(slot, bound._slots) + self.assertGreaterEqual(backend_slot.close_count, 1) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_pre_provider_interrupt_recycles_core_only(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + original_begin = TransferSlot._begin_generation_unlocked + + def interrupt_after_core_publication(self) -> None: + original_begin(self) + raise KeyboardInterrupt + + with mock.patch.object( + TransferSlot, + "_begin_generation_unlocked", + interrupt_after_core_publication, + ): + self.assertIsNone(slot.start()) + + self.assertTrue(slot.active) + self.assertEqual(slot.generation, 1) + self.assertEqual(backend_slot.start_count, 0) + self.assertEqual(slot.state, WorkState.FAILED) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 0) + + backend_slot.next_state = WorkState.COMPLETED + slot.start() + self.assertEqual(slot.generation, 2) + slot.recycle() + self.assertEqual(backend_slot.recycle_count, 1) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_recycle_interrupt_publishes_clean_idle_last(self) -> None: + class InterruptingIdleCommitTransferSlot(TransferSlot): + __slots__ = ("interrupt_idle_commit",) + + def __init__(self, *args, **kwargs) -> None: + self.interrupt_idle_commit = False + super().__init__(*args, **kwargs) + + def __setattr__(self, name, value) -> None: + object.__setattr__(self, name, value) + if ( + name == "_active" + and value is False + and getattr(self, "interrupt_idle_commit", False) + ): + object.__setattr__(self, "interrupt_idle_commit", False) + raise KeyboardInterrupt + + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + with mock.patch.object( + _api, "TransferSlot", InterruptingIdleCommitTransferSlot + ): + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + prior_error = RuntimeError("prior generation failed") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = prior_error + slot.start() + self.assertTrue(slot.test()) + self.assertIs(slot.error, prior_error) + + slot.interrupt_idle_commit = True + with self.assertRaises(KeyboardInterrupt): + slot.recycle() + + self.assertFalse(slot.active) + self.assertIsNone(slot.state) + self.assertIsNone(slot.error) + # The public idle commit precedes private cleanup. The retained + # terminal state is an explicit dirty-idle marker that the next begin + # must repair before publishing another generation. + self.assertEqual(slot._state, WorkState.FAILED) + self.assertTrue(slot._backend_generation_active) + self.assertIs(slot._error, prior_error) + + backend_slot.next_state = WorkState.COMPLETED + backend_slot.next_error = None + self.assertEqual(slot.execute(), 2) + self.assertIsNone(slot.error) + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_terminal_error_publication_retries_after_interrupt( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + backend_slot = source._backend.request_slots[0] + + failure = RuntimeError("retryable provider failure") + backend_slot.next_state = WorkState.FAILED + backend_slot.next_error = failure + backend_slot.raise_error_once = KeyboardInterrupt() + slot.start() + with self.assertRaises(KeyboardInterrupt): + _ = slot.state + self.assertEqual(slot._state, WorkState.PENDING) + self.assertEqual(slot.state, WorkState.FAILED) + self.assertIs(slot.error, failure) + slot.recycle() + + backend_slot.next_state = WorkState.CANCELLED + slot.start() + with ( + mock.patch( + "torch.distributed._transfer._api.TransferCancelledError", + side_effect=KeyboardInterrupt, + ), + self.assertRaises(KeyboardInterrupt), + ): + _ = slot.state + self.assertEqual(slot._state, WorkState.PENDING) + self.assertEqual(slot.state, WorkState.CANCELLED) + self.assertIsInstance(slot.error, TransferCancelledError) + slot.recycle() + + backend_slot.next_state = None + slot.start() + with ( + mock.patch( + "torch.distributed._transfer._api.BackendFailureError", + side_effect=KeyboardInterrupt, + ), + self.assertRaises(KeyboardInterrupt), + ): + _ = slot.state + self.assertEqual(slot._state, WorkState.PENDING) + self.assertEqual(slot.state, WorkState.FAILED) + self.assertIsInstance(slot.error, BackendFailureError) + slot.recycle() + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_transfer_slot_completion_event_is_capability_gated_and_generation_local( + self, + ) -> None: + source = Endpoint("source", backend="test-request-slot-no-completion-event") + destination = Endpoint( + "destination", backend="test-request-slot-no-completion-event" + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + slot = bound.create_slot() + + with self.assertRaisesRegex(BusyError, "no active generation"): + _ = slot.completion_event + slot.start() + self.assertIsNone(slot.completion_event) + slot.recycle() + with self.assertRaisesRegex(BusyError, "no active generation"): + _ = slot.completion_event + + slot.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_reference_retains_remote_owner_during_copy(self) -> None: + class WeakBytearray(bytearray): + pass + + source = Endpoint("source", thread_mode="multiple") + destination = Endpoint("destination", thread_mode="multiple") + source_registration = source.register(bytearray(b"data"), name="buffer") + destination_buffer = WeakBytearray(4) + destination_owner = weakref.ref(destination_buffer) + destination_registration = destination.register( + destination_buffer, name="buffer" + ) + del destination_buffer + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[source_registration.region()], + remote=[peer.region("buffer")], + ) + + entered_copy = threading.Event() + continue_copy = threading.Event() + original_copy = _reference._copy_spans + + def blocking_copy(source_spans, destination_spans, *, anchors=()): + entered_copy.set() + self.assertTrue(continue_copy.wait(1)) + self.assertIsNotNone(destination_owner()) + return original_copy(source_spans, destination_spans, anchors=anchors) + + _reference._copy_spans = blocking_copy + result = [] + writer = threading.Thread(target=lambda: result.append(source.write(plan))) + try: + writer.start() + self.assertTrue(entered_copy.wait(1)) + destination_registration.close() + gc.collect() + self.assertIsNotNone(destination_owner()) + continue_copy.set() + writer.join(1) + self.assertFalse(writer.is_alive()) + result[0].wait() + finally: + continue_copy.set() + writer.join(1) + _reference._copy_spans = original_copy + + result[0].close() + plan.close() + peer.close() + source_registration.close() + source.close() + destination.close() + + def test_reference_tensor_and_storage_registration(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_tensor = torch.arange(8, dtype=torch.uint8) + destination_tensor = torch.zeros(8, dtype=torch.uint8) + source_registration = source.register( + source_tensor.untyped_storage(), name="storage" + ) + destination_registration = destination.register( + destination_tensor, name="tensor" + ) + peer = destination.import_peer(bytes(source.export_metadata())) + plan = destination.prepare( + local=[destination_registration.region()], + remote=[peer.region("storage")], + ) + + work = destination.read(plan) + work.wait() + self.assertTrue(torch.equal(destination_tensor, source_tensor)) + + work.close() + plan.close() + peer.close() + destination_registration.close() + source_registration.close() + destination.close() + source.close() + + def test_tensor_rebinding_keeps_registered_storage_alive(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_tensor = torch.arange(8, dtype=torch.uint8) + expected = source_tensor.clone() + destination_tensor = torch.zeros_like(source_tensor) + source_registration = source.register(source_tensor, name="source") + destination_registration = destination.register( + destination_tensor, name="destination" + ) + source_tensor.set_(torch.full_like(source_tensor, 99)) + peer = destination.import_peer(bytes(source.export_metadata())) + plan = destination.prepare( + local=[destination_registration.region()], + remote=[peer.region("source")], + ) + + work = destination.read(plan) + work.wait() + self.assertTrue(torch.equal(destination_tensor, expected)) + + work.close() + plan.close() + peer.close() + destination_registration.close() + source_registration.close() + destination.close() + source.close() + + def test_indexed_plan_reuse(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"00112233") + destination_buffer = bytearray(b"----------") + source_registration = source.register(source_buffer, name="blocks") + destination_registration = destination.register( + destination_buffer, name="blocks" + ) + peer = source.import_peer(bytes(destination.export_metadata())) + local = [source_registration.region(index * 2, 2) for index in range(4)] + remote = [peer.region("blocks", index * 2, 2) for index in range(5)] + plan = source.prepare(local=local, remote=remote, indexed=True) + + first = source.write(plan, local_indices=[3, 1], remote_indices=[0, 2]) + self.assertTrue(first.test()) + self.assertEqual(destination_buffer, bytearray(b"33--11----")) + first.close() + second = source.write(plan, local_indices=[0], remote_indices=[4]) + second.wait() + self.assertEqual(destination_buffer, bytearray(b"33--11--00")) + + second.close() + plan.close() + peer.close() + source_registration.close() + destination_registration.close() + source.close() + destination.close() + + def test_index_selection_snapshots_packed_int32_and_reuses_dynamic_tag( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"a0b0c0d0") + destination_buffer = bytearray(b"--------") + local = source.register(source_buffer, name="source") + remote = destination.register(destination_buffer, name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=2, count=4)], + remote=[peer.region("destination", 0, 1, stride=2, count=4)], + indexed=True, + ) + local_indices = array("i", [3, 1]) + remote_indices = array("i", [0, 2]) + self.assertEqual(local_indices.itemsize, 4) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=local_indices, + remote_indices=remote_indices, + ) + self.assertIsInstance(selection, IndexSelection) + self.assertEqual(selection.local_count, 2) + self.assertEqual(selection.remote_count, 2) + self.assertEqual(selection.selected_nbytes, 2) + backend_selection = selection._backend_selection + self.assertIsInstance(backend_selection, BackendIndexSelection) + self.assertTrue(backend_selection.local_indices.readonly) + self.assertTrue(backend_selection.remote_indices.readonly) + self.assertIsInstance(backend_selection.local_indices.obj, bytes) + self.assertIs( + backend_selection.local_indices.obj, + backend_selection.remote_indices.obj, + ) + with self.assertRaisesRegex(AttributeError, "immutable"): + selection._op = TransferOp.READ + with self.assertRaisesRegex(AttributeError, "immutable"): + backend_selection._destination_indices_unique = False + + # The caller-owned mutable arrays no longer affect the snapshot. + local_indices[:] = array("i", [0, 0]) + remote_indices[:] = array("i", [1, 1]) + work = selection.submit(notification=b"first") + self.assertIn(selection, work._dependencies) + work.wait() + self.assertEqual(destination_buffer, bytearray(b"d---b---")) + self.assertEqual( + [item.payload for item in destination.poll_notifications()], + [b"first"], + ) + work.close() + + destination_buffer[:] = b"--------" + bound = selection.bind() + self.assertIs(bound._selection, selection) + work = bound.submit(notification=b"second") + work.wait() + self.assertEqual(destination_buffer, bytearray(b"d---b---")) + self.assertEqual( + [item.payload for item in destination.poll_notifications()], + [b"second"], + ) + work.close() + bound.close() + + immutable_bound = plan.bind( + TransferOp.WRITE, + local_indices=[3, 1], + remote_indices=[0, 2], + ) + with self.assertRaisesRegex(TypeError, "IndexSelection binding"): + immutable_bound.submit(notification=b"override") + immutable_bound.close() + plan.close() + with self.assertRaises(ClosedError): + selection.submit() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_fused_submit_and_bounded_work_poll(self) -> None: + source = Endpoint("source", backend="test-fused-selection") + destination = Endpoint("destination", backend="test-fused-selection") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + ) + backend = source._backend + self.assertIsInstance(backend, _FusedSelectionReferenceBackend) + + work, terminal = selection.submit_and_test( + notification=b"pending", + max_polls=17, + timeout_ns=23, + ) + self.assertFalse(terminal) + backend_work = backend.last_work + self.assertIsNotNone(backend_work) + assert backend_work is not None + self.assertEqual(backend_work.state_queries, 0) + self.assertEqual( + work.poll_bounded( + max_polls=31, + timeout_ns=37, + timeout_check_interval=7, + ), + WorkState.COMPLETED, + ) + self.assertEqual(backend_work.bounded_calls, [(31, 37, 7)]) + self.assertEqual(backend_work.state_queries, 0) + self.assertEqual(work.state, WorkState.COMPLETED) + work.close() + + backend.initial_state = WorkState.COMPLETED + work, state = selection.submit_and_poll(max_polls=0) + self.assertIs(state, WorkState.COMPLETED) + immediate_work = backend.last_work + assert immediate_work is not None + self.assertTrue(work.test()) + self.assertEqual(immediate_work.state_queries, 0) + self.assertEqual( + backend.fused_calls, + [(17, 23, b"pending"), (0, None, None)], + ) + + work.close() + work, terminal = selection.submit_and_test(max_polls=0) + self.assertTrue(terminal) + self.assertEqual(backend.last_work.state_queries, 0) + + for arguments, error_type in ( + ({"max_polls": True}, TypeError), + ({"max_polls": -1}, ValueError), + ({"timeout_ns": True}, TypeError), + ({"timeout_ns": -1}, ValueError), + ): + with self.assertRaises(error_type): + selection.submit_and_test(**arguments) + with self.assertRaises(TypeError): + work.test_bounded(max_polls=True) + with self.assertRaises(ValueError): + work.test_bounded(max_polls=0) + work.close() + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_plan_one_shot_indices_use_private_snapshot_and_fallback(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"a0b0c0d0") + destination_buffer = bytearray(b"--------") + local = source.register(source_buffer, name="source") + remote = destination.register(destination_buffer, name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=2, count=4)], + remote=[peer.region("destination", 0, 1, stride=2, count=4)], + indexed=True, + ) + local_indices = array("i", [3, 1]) + remote_indices = array("i", [0, 2]) + backend = source._backend + original_submit = backend.submit_prevalidated + captured: list[BackendIndexSelection] = [] + + def capture_submit(*args, **kwargs) -> None: + captured.append(kwargs["selection"]) + local_indices[:] = array("i", [0, 0]) + remote_indices[:] = array("i", [1, 1]) + original_submit(*args, **kwargs) + + backend.submit_prevalidated = capture_submit + work = plan.submit_indices( + TransferOp.WRITE, + local_indices=local_indices, + remote_indices=remote_indices, + notification=b"one-shot", + ) + self.assertEqual(destination_buffer, bytearray(b"d---b---")) + self.assertEqual(len(captured), 1) + backend_selection = captured[0] + self.assertIn(backend_selection, work._dependencies) + self.assertFalse( + any(isinstance(value, IndexSelection) for value in work._dependencies) + ) + self.assertTrue(backend_selection.local_indices.readonly) + self.assertTrue(backend_selection.remote_indices.readonly) + self.assertEqual(list(backend_selection.local_indices), [3, 1]) + self.assertEqual(list(backend_selection.remote_indices), [0, 2]) + self.assertEqual( + [item.payload for item in destination.poll_notifications()], + [b"one-shot"], + ) + work.close() + + destination_buffer[:] = b"--------" + work, terminal = plan.submit_indices_and_test( + TransferOp.WRITE, + local_indices=array("i", [2]), + remote_indices=array("i", [3]), + max_polls=0, + ) + # The source-compatible provider cannot expose its post-derived state + # without a poll, so zero remains conservatively nonterminal. + self.assertFalse(terminal) + self.assertTrue(work.test()) + self.assertEqual(destination_buffer, bytearray(b"------c-")) + work.close() + + destination_buffer[:] = b"--------" + work, state = plan.submit_indices_and_poll( + TransferOp.WRITE, + local_indices=array("i", [1]), + remote_indices=array("i", [1]), + max_polls=1, + timeout_ns=None, + ) + self.assertIs(state, WorkState.COMPLETED) + self.assertEqual(destination_buffer, bytearray(b"--b-----")) + work.close() + + # caller_ready never consumes local device identities. In particular, + # a nonuniform-device catalog must not trigger an O(k) region walk just + # to build facts that this ordering mode discards. + plan._local_uniform_device = None + with mock.patch.object( + _api, + "_select_region_blocks", + side_effect=AssertionError("caller_ready walked local devices"), + ): + work = plan.submit_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + ) + work.close() + + with self.assertRaisesRegex(ValueError, "non-empty"): + plan.submit_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + notification=b"", + ) + with self.assertRaisesRegex(ValueError, "does not accept"): + plan.submit_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + stream=object(), + ) + for method_name in ( + "submit_indices", + "submit_indices_and_poll", + "submit_indices_and_test", + ): + method = getattr(plan, method_name) + with self.subTest(method=method_name, argument="op"): + with self.assertRaisesRegex(ValueError, "op must be one of"): + method( + None, + local_indices=[0], + remote_indices=[0], + ) + with self.subTest(method=method_name, argument="local_indices"): + with self.assertRaisesRegex( + InvalidRegionError, "local_indices must be non-empty" + ): + method( + TransferOp.WRITE, + local_indices=None, + remote_indices=[0], + ) + with self.subTest(method=method_name, argument="remote_indices"): + with self.assertRaisesRegex( + InvalidRegionError, "remote_indices must be non-empty" + ): + method( + TransferOp.WRITE, + local_indices=[0], + remote_indices=None, + ) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_plan_one_shot_indices_enter_fused_provider_boundary(self) -> None: + source = Endpoint("source", backend="test-fused-selection") + destination = Endpoint("destination", backend="test-fused-selection") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + backend = source._backend + self.assertIsInstance(backend, _FusedSelectionReferenceBackend) + backend.initial_state = WorkState.COMPLETED + + work, state = plan.submit_indices_and_poll( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + notification=b"fused-one-shot", + max_polls=0, + timeout_ns=None, + ) + self.assertIs(state, WorkState.COMPLETED) + self.assertEqual(backend.fused_calls, [(0, None, b"fused-one-shot")]) + self.assertTrue( + any( + isinstance(value, BackendIndexSelection) for value in work._dependencies + ) + ) + self.assertFalse( + any(isinstance(value, IndexSelection) for value in work._dependencies) + ) + work.close() + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_fused_selection_immediate_state_retains_epilogue_interrupt(self) -> None: + class InterruptOnceOnDiscard(set): + armed = True + + def discard(self, value) -> None: + super().discard(value) + if self.armed: + self.armed = False + raise KeyboardInterrupt + + source = Endpoint("source", backend="test-fused-selection") + destination = Endpoint("destination", backend="test-fused-selection") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + ) + backend = source._backend + self.assertIsInstance(backend, _FusedSelectionReferenceBackend) + backend.initial_state = WorkState.COMPLETED + plan._submitting = InterruptOnceOnDiscard() + + work, terminal = selection.submit_and_test(max_polls=0) + self.assertTrue(terminal) + self.assertFalse(plan._submitting) + self.assertEqual(source.open_works, (work,)) + self.assertEqual(work.state, WorkState.FAILED) + self.assertIsInstance(work.error, KeyboardInterrupt) + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_fused_selection_pending_state_reconciles_epilogue_interrupt(self) -> None: + class InterruptOnceOnDiscard(set): + armed = True + + def discard(self, value) -> None: + super().discard(value) + if self.armed: + self.armed = False + raise KeyboardInterrupt + + source = Endpoint("source", backend="test-fused-selection") + destination = Endpoint("destination", backend="test-fused-selection") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + ) + backend = source._backend + self.assertIsInstance(backend, _FusedSelectionReferenceBackend) + backend.initial_state = WorkState.PENDING + plan._submitting = InterruptOnceOnDiscard() + + work, terminal = selection.submit_and_test(max_polls=0) + self.assertFalse(terminal) + self.assertFalse(plan._submitting) + self.assertEqual(source.open_works, (work,)) + self.assertTrue(work.test_bounded(max_polls=1, timeout_ns=None)) + self.assertEqual(work.state, WorkState.FAILED) + self.assertIsInstance(work.error, KeyboardInterrupt) + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_accepts_index_scalars_and_rejects_unsafe_values( + self, + ) -> None: + class IndexScalar: + def __init__(self, value: int) -> None: + self.value = value + + def __index__(self) -> int: + return self.value + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + + selection = plan.select_indices( + "write", + local_indices=[IndexScalar(3), IndexScalar(1)], + remote_indices=[IndexScalar(0), IndexScalar(2)], + ) + work = selection.submit() + work.close() + for local_indices, remote_indices, message in ( + ([True], [0], "integers"), + ([0], [False], "integers"), + ([4], [0], "outside"), + ([0, 1], [2, 2], "overlap"), + ([], [], "non-empty"), + ): + with self.assertRaisesRegex(InvalidRegionError, message): + plan.select_indices( + TransferOp.WRITE, + local_indices=local_indices, + remote_indices=remote_indices, + ) + + # The raw convenience fallback now follows the same operator.index + # contract, without requiring callers to construct a selection. + work = source.write( + plan, + local_indices=[IndexScalar(0)], + remote_indices=[IndexScalar(1)], + ) + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_native_int32_analyzer_matches_fallback(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + block_count = _api._INDEX_SELECTION_TORCH_THRESHOLD + 3 + local = source.register(bytearray(block_count), name="source") + remote = destination.register(bytearray(block_count), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=block_count)], + remote=[peer.region("destination", 0, 1, stride=1, count=block_count)], + indexed=True, + ) + cutoff = _api._INDEX_SELECTION_TORCH_THRESHOLD + signatures = [] + for analyzer in (None, _emulated_compiled_index_analyzer): + with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", range(cutoff + 1)), + remote_indices=array("i", reversed(range(cutoff + 1))), + ) + backend_selection = selection._backend_selection + signatures.append( + ( + bytes(backend_selection.local_indices), + bytes(backend_selection.remote_indices), + selection.local_count, + selection.remote_count, + selection.selected_nbytes, + ) + ) + self.assertEqual(signatures[0], signatures[1]) + + caller_local = array("i", range(cutoff + 1)) + caller_remote = array("i", reversed(range(cutoff + 1))) + + def mutate_after_snapshot(*args, **kwargs): + result = _emulated_compiled_index_analyzer(*args, **kwargs) + caller_local[:] = array("i", [0]) * len(caller_local) + caller_remote[:] = array("i", [0]) * len(caller_remote) + return result + + analyzer = mock.Mock(side_effect=mutate_after_snapshot) + with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=caller_local, + remote_indices=caller_remote, + ) + analyzer.assert_called_once() + self.assertEqual( + list(selection._backend_selection.local_indices), + list(range(cutoff + 1)), + ) + self.assertEqual( + list(selection._backend_selection.remote_indices), + list(reversed(range(cutoff + 1))), + ) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_compiled_analyzer_preserves_error_precedence( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(8), name="source") + remote = destination.register(bytearray(8), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=8)], + remote=[peer.region("destination", 0, 1, stride=1, count=8)], + indexed=True, + ) + cases = ( + (array("i", [-1]), array("i", [8]), "local_indices.*outside"), + (array("i", [0]), array("i", [8]), "remote_indices.*outside"), + ( + array("i", [0, 1, 2]), + array("i", [0, 0]), + "byte totals differ", + ), + (array("i", [0, 1]), array("i", [0, 0]), "overlap"), + ) + for analyzer in (None, _emulated_compiled_index_analyzer): + for local_indices, remote_indices, message in cases: + with ( + self.subTest(analyzer=analyzer, message=message), + mock.patch.object(_api, "_compiled_index_analyzer", analyzer), + self.assertRaisesRegex(InvalidRegionError, message), + ): + plan.select_indices( + TransferOp.WRITE, + local_indices=local_indices, + remote_indices=remote_indices, + ) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_mixed_inputs_preserve_error_order_and_snapshot( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(4), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + + cases = ( + (array("i", [-1]), [True], r"local_indices index -1.*outside"), + (array("i", [4]), [False], r"local_indices index 4.*outside"), + (array("i", [-1]), [4], r"local_indices index -1.*outside"), + (array("i", [-1]), [], r"local_indices index -1.*outside"), + ([True], array("i", [4]), r"local_indices must contain integers"), + ([4], array("i", [-1]), r"local_indices index 4.*outside"), + ([], array("i", [-1]), r"local_indices must be non-empty"), + ) + for local_indices, remote_indices, message in cases: + with ( + self.subTest( + local_type=type(local_indices).__name__, + remote_type=type(remote_indices).__name__, + ), + self.assertRaisesRegex(InvalidRegionError, message), + ): + plan.select_indices( + TransferOp.WRITE, + local_indices=local_indices, + remote_indices=remote_indices, + ) + + caller_local = array("i", [3, 1]) + + class MutatingRemote: + def __iter__(self): + caller_local[:] = array("i", [0, 0]) + return iter((0, 2)) + + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=caller_local, + remote_indices=MutatingRemote(), + ) + self.assertEqual(list(selection._backend_selection.local_indices), [3, 1]) + self.assertEqual(list(selection._backend_selection.remote_indices), [0, 2]) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_compiled_analyzer_exception_is_not_retried( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(4), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + analyzer = mock.Mock(side_effect=RuntimeError("compiled helper failed")) + with ( + mock.patch.object(_api, "_compiled_index_analyzer", analyzer), + self.assertRaisesRegex(RuntimeError, "compiled helper failed"), + ): + plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [0, 1]), + remote_indices=array("i", [2, 3]), + ) + analyzer.assert_called_once() + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_missing_helper_and_third_party_sequences(self) -> None: + class IndexScalar: + def __init__(self, value: int) -> None: + self.value = value + + def __index__(self) -> int: + return self.value + + class ThirdPartyIndices: + def __init__(self, *values: int) -> None: + self.values = values + + def __iter__(self): + return iter(IndexScalar(value) for value in self.values) + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + + with mock.patch.object(_api, "_compiled_index_analyzer", None): + fallback = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + ) + self.assertEqual(list(fallback._backend_selection.local_indices), [3, 1]) + + analyzer = mock.Mock(side_effect=AssertionError("must not be called")) + with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): + mixed = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [2, 0]), + remote_indices=ThirdPartyIndices(1, 3), + ) + analyzer.assert_not_called() + self.assertEqual(list(mixed._backend_selection.remote_indices), [1, 3]) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_compiled_index_analyzer_defensive_buffer_contract(self) -> None: + analyzer = _api._compiled_index_analyzer + if analyzer is None: + self.skipTest("requires torch._C rebuilt with the transfer helper") + + local = array("i", [3, 1, 2]) + remote = array("i", [0, 0, 2]) + result = analyzer(local, remote, True) + self.assertEqual(result[0], local.tobytes() + remote.tobytes()) + self.assertEqual(result[1:7], (3, 3, 1, 3, 0, 2)) + self.assertTrue(result[7]) + + targeted_destinations = ( + array("i", [5]), + array("i", [5, 5]), + array("i", range(4096)), + array("i", reversed(range(4096))), + array("i", [2, 0, 3, 1]), + array("i", [7, 7, -5, 100]), + array("i", ((index * 2053) % 4096 for index in range(4096))), + array( + "i", + (((index * 37) % 64) * 104729 for index in range(64)), + ), + array("i", [0, 104729, -65537, 209458]), + array("i", [7, 8, 7]), + array("i", [0, 1, 2, 3, 4, 5, 6, 7, -1, 9, 10]), + array("i", [*range(4095), -1]), + array("i", [*reversed(range(1, 4096)), 4096]), + array("i", [-(2**31), 2**31 - 1, 0]), + ) + for destination in targeted_destinations: + expected_duplicate = len(set(destination)) != len(destination) + for destination_is_remote in (False, True): + with self.subTest( + destination_is_remote=destination_is_remote, + destination=list(destination[:8]), + ): + if destination_is_remote: + candidate = analyzer(local, destination, True) + else: + candidate = analyzer(destination, remote, False) + self.assertEqual(candidate[7], expected_duplicate) + destination_bounds = ( + candidate[5:7] if destination_is_remote else candidate[3:5] + ) + self.assertEqual( + destination_bounds, (min(destination), max(destination)) + ) + + invalid_buffers = ( + array("i"), + array("I", [0]), + array("q", [0]), + memoryview(array("i", [0, 1, 2, 3]))[::2], + ) + for invalid in invalid_buffers: + with self.subTest(invalid=invalid), self.assertRaises(ValueError): + analyzer(invalid, array("i", [0]), True) + + def test_index_selection_reuse_token_is_lazy_and_not_redundant(self) -> None: + source = Endpoint("source", backend="test-request-slot") + destination = Endpoint("destination", backend="test-request-slot") + local = source.register(bytearray(b"abcd"), name="source") + remote = destination.register(bytearray(4), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("destination", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [3, 1]), + remote_indices=array("i", [0, 2]), + ) + backend = source._backend + original_submit = backend.submit_prevalidated + observed_tokens: list[object | None] = [] + + def capture_submit(*args, **kwargs) -> None: + observed_tokens.append(kwargs["reuse_token"]) + original_submit(*args, **kwargs) + + backend.submit_prevalidated = capture_submit + self.assertIsNone(selection._reuse_token) + direct = selection.submit() + direct.close() + self.assertEqual(observed_tokens, [None]) + self.assertIsNone(selection._reuse_token) + + first_bound = selection.bind() + self.assertIsNone(first_bound._reuse_token) + slot = first_bound.create_slot() + self.assertIsNone(selection._reuse_token) + slot.close() + + first = first_bound.submit() + token = selection._reuse_token + self.assertIsNotNone(token) + first.close() + second = first_bound.submit() + second.close() + second_bound = selection.bind() + self.assertIsNone(second_bound._reuse_token) + third = second_bound.submit() + third.close() + self.assertEqual(len(observed_tokens), 4) + self.assertIsNone(observed_tokens[0]) + self.assertIs(observed_tokens[1], token) + self.assertIs(observed_tokens[2], token) + self.assertIs(observed_tokens[3], token) + + second_bound.close() + first_bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_index_selection_irregular_layout_keeps_conservative_fallback( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"abc"), name="source") + remote = destination.register(bytearray(3), name="destination") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1), local.region(1, 2)], + remote=[peer.region("destination", 0, 3)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[0, 1], + remote_indices=[0], + ) + self.assertFalse(selection._backend_selection.uniform_pair_layout) + self.assertEqual(selection.selected_nbytes, 3) + work = selection.submit() + work.wait() + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_backend_index_selection_cannot_be_forged(self) -> None: + with self.assertRaisesRegex(TypeError, "created by Core"): + BackendIndexSelection( + TransferOp.WRITE, + object(), + memoryview(array("i", [0])), + memoryview(array("i", [0])), + selected_nbytes=1, + uniform_pair_layout=True, + destination_catalog_nonoverlapping=True, + destination_indices_unique=True, + _token=object(), + ) + + def test_indexed_compact_strided_block_catalog(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"a0b0c0d0") + destination_buffer = bytearray(b"--------") + source_registration = source.register(source_buffer, name="source") + destination_registration = destination.register( + destination_buffer, name="destination" + ) + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[source_registration.region(0, 1, stride=2, count=4)], + remote=[peer.region("destination", 0, 1, stride=2, count=4)], + indexed=True, + ) + + self.assertEqual(plan.local_block_count, 4) + self.assertEqual(plan.remote_block_count, 4) + work = source.write( + plan, + local_indices=[3, 1], + remote_indices=[0, 2], + ) + work.wait() + self.assertEqual(destination_buffer, bytearray(b"d---b---")) + work.close() + with self.assertRaises(InvalidRegionError): + source.write(plan, local_indices=[4], remote_indices=[0]) + + plan.close() + peer.close() + source_registration.close() + destination_registration.close() + source.close() + destination.close() + + def test_indexed_logical_block_sizes_must_match(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + destination_buffer = bytearray(4) + source_registration = source.register(bytearray(b"a0b0"), name="source") + destination_registration = destination.register( + destination_buffer, name="destination" + ) + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[source_registration.region(0, 1, stride=2, count=2)], + remote=[peer.region("destination", 0, 2, stride=2, count=2)], + indexed=True, + ) + + with self.assertRaisesRegex(InvalidRegionError, "byte totals differ"): + source.write(plan, local_indices=[0], remote_indices=[0]) + work = source.write( + plan, + local_indices=[0, 1], + remote_indices=[1], + ) + work.wait() + self.assertEqual(destination_buffer, bytearray(b"\x00\x00ab")) + work.close() + + plan.close() + peer.close() + source_registration.close() + destination_registration.close() + source.close() + destination.close() + + def test_strided_regions(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"a0b0c0d0") + destination_buffer = bytearray(4) + source_registration = source.register(source_buffer, name="source") + destination_registration = destination.register( + destination_buffer, name="destination" + ) + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[source_registration.region(0, 1, stride=2, count=4)], + remote=[peer.region("destination")], + ) + + work = source.write(plan) + work.wait() + self.assertEqual(destination_buffer, bytearray(b"abcd")) + + work.close() + plan.close() + peer.close() + source_registration.close() + destination_registration.close() + source.close() + destination.close() + + def test_notifications_identify_source(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + + with self.assertRaisesRegex(ValueError, "non-empty"): + source.write(plan, notification=b"") + transfer = source.write(plan, notification=b"visible") + transfer.wait() + notices = destination.poll_notifications() + self.assertEqual(len(notices), 1) + self.assertEqual(notices[0].payload, b"visible") + self.assertEqual(notices[0].source_endpoint_id, source.endpoint_id) + self.assertEqual(notices[0].source_incarnation, source.incarnation) + self.assertEqual(notices[0].source, source.identity) + + class OversizedPayload(bytearray): + converted = False + + def __bytes__(self) -> bytes: + self.converted = True + return super().__bytes__() + + oversized = OversizedPayload(b"1234") + with ( + mock.patch("torch.distributed._transfer._api._MAX_NOTIFICATION_BYTES", 3), + self.assertRaisesRegex(ValueError, "size limit"), + ): + source.send_notification(peer, oversized) + self.assertFalse(oversized.converted) + + standalone = source.send_notification(peer, b"standalone") + standalone.wait() + self.assertEqual( + [item.payload for item in destination.poll_notifications()], + [b"standalone"], + ) + + self.assertTrue(source.capabilities.synchronous_notification_send) + self.assertEqual(set(source.open_works), {standalone, transfer}) + sync_payload = b"sync-exact-payload" + self.assertIsNone(source.send_notification_sync(peer, sync_payload)) + self.assertEqual(set(source.open_works), {standalone, transfer}) + sync_notifications = destination.poll_notifications() + self.assertEqual([item.payload for item in sync_notifications], [sync_payload]) + self.assertIs(sync_notifications[0].payload, sync_payload) + + provider_failure = RuntimeError("provider failed after entry") + + def fail_after_entry(*args, **kwargs) -> None: + del args, kwargs + raise provider_failure + + source._send_notification_sync_dispatch = fail_after_entry + with self.assertRaisesRegex( + BackendFailureError, "delivery outcome may be ambiguous" + ) as context: + source.send_notification_sync(peer, b"ambiguous") + self.assertIs(context.exception.__cause__, provider_failure) + + transfer_error = ClosedError("provider is closed") + + def fail_with_transfer_error(*args, **kwargs) -> None: + del args, kwargs + raise transfer_error + + source._send_notification_sync_dispatch = fail_with_transfer_error + with self.assertRaises(ClosedError) as context: + source.send_notification_sync(peer, b"rejected") + self.assertIs(context.exception, transfer_error) + + source._send_notification_sync_dispatch = lambda *args, **kwargs: object() + with self.assertRaisesRegex( + BackendFailureError, "must return None.*outcome may be ambiguous" + ): + source.send_notification_sync(peer, b"invalid-result") + self.assertEqual(set(source.open_works), {standalone, transfer}) + + standalone.close() + transfer.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_attached_only_backend_can_poll_notifications(self) -> None: + endpoint = Endpoint("endpoint", backend="test-attached-only") + destination = Endpoint("destination", backend="test-attached-only") + peer = endpoint.import_peer(bytes(destination.export_metadata())) + self.assertTrue(endpoint.capabilities.attached_notifications) + self.assertFalse(endpoint.capabilities.standalone_notifications) + self.assertFalse(endpoint.capabilities.synchronous_notification_send) + with self.assertRaisesRegex(UnsupportedError, "synchronous notification"): + endpoint.send_notification_sync(peer, b"unsupported") + with self.assertRaisesRegex(UnsupportedError, "synchronous standalone"): + _api.Backend.send_notification_sync( + endpoint._backend, object(), b"unsupported" + ) + self.assertEqual(endpoint.poll_notifications(), []) + peer.close() + endpoint.close() + destination.close() + + def test_multiple_synchronous_notification_sends_overlap_and_pin_peer( + self, + ) -> None: + source = Endpoint("source", thread_mode="multiple") + destination = Endpoint("destination", thread_mode="multiple") + peer = source.import_peer(bytes(destination.export_metadata())) + original_send = source._send_notification_sync_dispatch + self.assertIsNotNone(original_send) + entered_first = threading.Event() + entered_both = threading.Event() + release = threading.Event() + count_lock = threading.Lock() + entered = 0 + errors: list[BaseException] = [] + + def blocking_send(peer_handle, payload) -> None: + nonlocal entered + with count_lock: + entered += 1 + entered_first.set() + if entered == 2: + entered_both.set() + if not release.wait(2): + raise RuntimeError("timed out waiting to release concurrent sends") + assert original_send is not None + original_send(peer_handle, payload) + + source._send_notification_sync_dispatch = blocking_send + + def send(payload: bytes) -> None: + try: + source.send_notification_sync(peer, payload) + except BaseException as error: + errors.append(error) + + threads = [ + threading.Thread(target=send, args=(payload,)) + for payload in (b"first", b"second") + ] + try: + for thread in threads: + thread.start() + self.assertTrue(entered_first.wait(1)) + self.assertTrue( + entered_both.wait(1), + "MULTIPLE synchronous sends were serialized before native entry", + ) + with self.assertRaisesRegex(BusyError, "synchronous notification"): + peer.close() + finally: + release.set() + for thread in threads: + thread.join(2) + + self.assertTrue(all(not thread.is_alive() for thread in threads)) + self.assertEqual(errors, []) + self.assertEqual( + sorted(item.payload for item in destination.poll_notifications()), + [b"first", b"second"], + ) + peer.close() + source.close() + destination.close() + + def test_multiple_synchronous_notification_interrupted_cleanup_cannot_wedge_peer( + self, + ) -> None: + class InterruptSecondEnter: + def __init__(self, lock) -> None: + self.lock = lock + self.enters = 0 + + def __enter__(self): + self.enters += 1 + if self.enters == 2: + raise KeyboardInterrupt + return self.lock.__enter__() + + def __exit__(self, *args): + return self.lock.__exit__(*args) + + source = Endpoint("source", thread_mode="multiple") + destination = Endpoint("destination", thread_mode="multiple") + peer = source.import_peer(bytes(destination.export_metadata())) + original_lock = source._state_lock + source._state_lock = InterruptSecondEnter(original_lock) + retained_error = None + try: + try: + source.send_notification_sync(peer, b"accepted-before-interrupt") + except KeyboardInterrupt as error: + # Retain the traceback, and therefore the completed send frame, + # through peer.close(). Close must use the explicit completion + # bit rather than mistake weakref liveness for native activity. + retained_error = error + else: + self.fail("cleanup-lock interruption did not propagate") + finally: + source._state_lock = original_lock + + gc.collect() + self.assertIsNotNone(retained_error) + self.assertTrue(peer._active_sync_notification_sends) + self.assertTrue( + all( + reference() is None or not reference().active + for reference in peer._active_sync_notification_sends + ) + ) + # close() prunes dead reservations and cannot mistake interrupted + # cleanup for a still-running native call. + peer.close() + self.assertFalse(peer._active_sync_notification_sends) + retained_error = None + source.close() + destination.close() + + def test_multiple_async_notification_send_overlaps_manual_progress(self) -> None: + # This is a backend-entry concurrency test, not a timing benchmark. + source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) + destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) + peer = source.import_peer(bytes(destination.export_metadata())) + backend_work = _PendingWork() + send_entered = threading.Event() + progress_entered = threading.Event() + release_send = threading.Event() + errors: list[BaseException] = [] + returned: list[Work] = [] + + def blocked_send(*args, **kwargs) -> None: + kwargs["adopt_work"](backend_work) + send_entered.set() + if not release_send.wait(2): + raise RuntimeError("test did not release notification send") + + def observe_progress() -> None: + progress_entered.set() + + source._backend.send_notification = blocked_send + source._backend.progress = observe_progress + + def send() -> None: + try: + returned.append(source.send_notification(peer, b"overlap")) + except BaseException as error: + errors.append(error) + + send_thread = threading.Thread(target=send) + send_thread.start() + try: + self.assertTrue(send_entered.wait(1)) + progress_thread = threading.Thread(target=source.progress) + progress_thread.start() + progress_thread.join(1) + self.assertFalse(progress_thread.is_alive()) + self.assertTrue(progress_entered.is_set()) + with self.assertRaises(BusyError): + peer.close() + finally: + release_send.set() + send_thread.join(2) + + self.assertFalse(send_thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(len(returned), 1) + backend_work._state = WorkState.COMPLETED + returned[0].close() + peer.close() + source.close() + destination.close() + + def test_multiple_notification_consumer_admission_serializes_and_blocks_close( + self, + ) -> None: + endpoint = Endpoint("consumer", thread_mode=ThreadMode.MULTIPLE) + first_entered = threading.Event() + second_entered = threading.Event() + release_first = threading.Event() + count_lock = threading.Lock() + entries = 0 + errors: list[BaseException] = [] + + def blocked_poll(max_items=None): + del max_items + nonlocal entries + with count_lock: + entries += 1 + if entries == 1: + first_entered.set() + else: + second_entered.set() + if entries == 1 and not release_first.wait(2): + raise RuntimeError("test did not release first notification poll") + return [] + + endpoint._backend.poll_notifications = blocked_poll + + def poll() -> None: + try: + endpoint.poll_notifications() + except BaseException as error: + errors.append(error) + + first = threading.Thread(target=poll) + second = threading.Thread(target=poll) + first.start() + self.assertTrue(first_entered.wait(1)) + second.start() + self.assertFalse(second_entered.wait(0.1)) + + close_errors: list[BaseException] = [] + + def close() -> None: + try: + endpoint.close() + except BaseException as error: + close_errors.append(error) + + close_thread = threading.Thread(target=close) + close_thread.start() + close_thread.join(1) + self.assertFalse(close_thread.is_alive()) + self.assertEqual(len(close_errors), 1) + self.assertIsInstance(close_errors[0], BusyError) + self.assertFalse(endpoint.closed) + + release_first.set() + first.join(2) + second.join(2) + self.assertTrue(second_entered.is_set()) + self.assertTrue(all(not thread.is_alive() for thread in (first, second))) + self.assertEqual(errors, []) + (warmed_reservation,) = endpoint._multiple_backend_call_reservation_pool + endpoint.poll_notifications() + self.assertEqual(len(endpoint._multiple_backend_call_reservation_pool), 1) + self.assertIs( + endpoint._multiple_backend_call_reservation_pool[0], warmed_reservation + ) + endpoint.close() + + def test_notification_batches_are_immutable_bounded_and_source_compatible( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + peer = source.import_peer(bytes(destination.export_metadata())) + first = source.send_notification(peer, b"first") + second_payload = bytearray(b"second") + second = source.send_notification(peer, second_payload) + first.wait() + second.wait() + second_payload[:] = b"mutate" + + with mock.patch.object( + destination._backend, + "poll_notification_batches", + side_effect=AssertionError("zero max_items must not enter the backend"), + ): + self.assertEqual(destination.poll_notification_batches(max_items=0), []) + for invalid, error_type in ( + (True, TypeError), + (-1, ValueError), + (1.0, TypeError), + ): + with self.assertRaises(error_type): + destination.poll_notification_batches(max_items=invalid) # type: ignore[arg-type] + + batches = destination.poll_notification_batches() + self.assertEqual(len(batches), 1) + self.assertIsInstance(batches[0], NotificationBatch) + self.assertEqual(batches[0].payloads, (b"first", b"second")) + self.assertEqual(batches[0].source, source.identity) + + exact_payloads = (b"retained-one", b"retained-two") + retained = BackendNotificationBatch( + exact_payloads, + source_endpoint_id=source.endpoint_id, + source_incarnation=source.incarnation, + ) + self.assertIs(retained.payloads, exact_payloads) + self.assertEqual(retained.total_nbytes, sum(map(len, exact_payloads))) + with ( + mock.patch.object( + destination._backend, + "poll_notification_batches", + return_value=[retained], + ), + mock.patch.object( + sys.modules[Endpoint.__module__], + "NotificationBatch", + side_effect=lambda *args, **kwargs: ( + self.assertFalse(destination._state_lock._is_owned()) + or NotificationBatch(*args, **kwargs) + ), + ), + ): + public = destination.poll_notification_batches() + self.assertIs(public[0].payloads, exact_payloads) + + mutable_payload = bytearray(b"snapshotted") + snapshotted = BackendNotificationBatch([mutable_payload]) + mutable_payload[:] = b"overwritten!" + self.assertEqual(snapshotted.payloads, (b"snapshotted",)) + self.assertEqual(snapshotted.total_nbytes, len(b"snapshotted")) + + class OversizedBytearray(bytearray): + def __bytes__(self) -> bytes: + raise AssertionError("oversized mutable input must not be copied") + + oversized_mutable = OversizedBytearray( + b"x" * (BackendNotificationBatch.MAX_PAYLOAD_BYTES + 1) + ) + with self.assertRaisesRegex(ValueError, "wire size limit"): + BackendNotificationBatch((oversized_mutable,)) + + overrun = BackendNotificationBatch( + (b"one", b"two"), + source_endpoint_id=source.endpoint_id, + source_incarnation=source.incarnation, + ) + with ( + mock.patch.object( + destination._backend, + "poll_notification_batches", + return_value=[overrun], + ), + self.assertRaisesRegex(BackendFailureError, "more notification payloads"), + ): + destination.poll_notification_batches(max_items=1) + + with self.assertRaisesRegex(ValueError, "at least one payload"): + BackendNotificationBatch(()) + with self.assertRaisesRegex(ValueError, "specified together"): + NotificationBatch((b"payload",), source_endpoint_id="incomplete") + + first.close() + second.close() + peer.close() + source.close() + destination.close() + + def test_notification_send_uses_work_adoption_and_direct_dispatch(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + peer = source.import_peer(bytes(destination.export_metadata())) + backend_send = source._backend.send_notification + observed_adoption = None + + def capture_send(*args, **kwargs) -> None: + nonlocal observed_adoption + observed_adoption = kwargs["adopt_work"] + backend_send(*args, **kwargs) + + source._backend.send_notification = capture_send + with mock.patch.object( + source, + "_backend_call", + side_effect=AssertionError( + "standalone notification must use its direct backend helper" + ), + ): + work = source.send_notification(peer, b"direct") + + self.assertIs(observed_adoption.__self__, work) + self.assertIs(observed_adoption.__func__, Work._adopt_backend_work) + work.wait() + work.close() + peer.close() + source.close() + destination.close() + + def test_notification_interruption_after_adoption_retains_peer(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + peer = source.import_peer(bytes(destination.export_metadata())) + backend_work = _PendingWork() + + def adopt_then_interrupt(*args, **kwargs) -> None: + kwargs["adopt_work"](backend_work) + raise KeyboardInterrupt + + source._backend.send_notification = adopt_then_interrupt + recovery_token = object() + work = source.send_notification( + peer, + b"message", + recovery_token=recovery_token, + ) + + self.assertEqual(source.open_works, (work,)) + self.assertIs(work.recovery_token, recovery_token) + with self.assertRaises(BusyError): + peer.close() + backend_work._state = WorkState.COMPLETED + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + + work.close() + self.assertIsNone(work.recovery_token) + peer.close() + source.close() + destination.close() + + def test_lifetime_order_is_enforced(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + + with self.assertRaises(BusyError): + local.close() + with self.assertRaises(BusyError): + peer.close() + with self.assertRaises(BusyError): + plan.close() + with self.assertRaises(BusyError): + source.close() + + work.close() + self.assertIsNone(work._plan) + work.close() + plan.close() + self.assertFalse(plan._local) + self.assertFalse(plan._remote) + plan.close() + peer.close() + self.assertIsNone(peer._handle) + peer.close() + local.close() + self.assertIsNone(local._memory.owner) + local.close() + remote.close() + source.close() + source.close() + destination.close() + + def test_plan_cannot_close_during_submit_bookkeeping(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + submit_entered = threading.Event() + allow_submit = threading.Event() + close_done = threading.Event() + close_errors: list[BaseException] = [] + original_submit = source._backend.submit + + def blocked_submit(*args, **kwargs): + submit_entered.set() + allow_submit.wait() + return original_submit(*args, **kwargs) + + source._backend.submit = blocked_submit + work_items: list[Work] = [] + submit_thread = threading.Thread( + target=lambda: work_items.append(source.write(plan)) + ) + + def close_plan() -> None: + try: + plan.close() + except BaseException as error: + close_errors.append(error) + finally: + close_done.set() + + submit_thread.start() + self.assertTrue(submit_entered.wait(1)) + close_thread = threading.Thread(target=close_plan) + close_thread.start() + self.assertTrue(close_done.wait(1)) + self.assertEqual(len(close_errors), 1) + self.assertIsInstance(close_errors[0], BusyError) + self.assertFalse(plan.closed) + allow_submit.set() + submit_thread.join() + close_thread.join() + self.assertEqual(len(work_items), 1) + self.assertFalse(plan.closed) + + work_items[0].close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_overlapping_transfer_is_request_local_failure(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + shared = bytearray(b"data") + local = source.register(shared, name="local") + remote = destination.register(shared, name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + + work = source.write(plan) + self.assertEqual(work.state, WorkState.FAILED) + self.assertIsInstance(work.error, InvalidRegionError) + with self.assertRaises(InvalidRegionError): + work.wait() + + work.close() + self.assertIsNone(work._error) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_indexed_destination_must_not_overlap(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + destination_buffer = bytearray(b"---") + local = source.register(bytearray(b"ab"), name="buffer") + remote = destination.register(destination_buffer, name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=2)], + remote=[ + peer.region("buffer", 0, 1, stride=2, count=2), + peer.region("buffer", 1, 1, stride=1, count=2), + ], + indexed=True, + ) + + work = source.write( + plan, + local_indices=[0, 1], + remote_indices=[0, 2], + ) + work.wait() + self.assertEqual(destination_buffer, bytearray(b"ab-")) + work.close() + + with self.assertRaises(InvalidRegionError): + source.write( + plan, + local_indices=[0, 1], + remote_indices=[1, 3], + ) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_overlapping_registrations_are_rejected(self) -> None: + endpoint = Endpoint("endpoint") + storage = bytearray(b"abcd") + registration = endpoint.register(storage, name="first") + with self.assertRaisesRegex( + InvalidRegionError, "overlaps existing registration" + ): + endpoint.register(storage, name="alias") + registration.close() + endpoint.close() + + def test_reference_deregister_is_idempotent_after_success(self) -> None: + endpoint = Endpoint("endpoint") + registration = endpoint.register(bytearray(4), name="buffer") + handle = registration._handle + + endpoint._backend.deregister(handle) + endpoint._backend.deregister(handle) + registration.close() + endpoint.close() + + def test_one_shot_plan_is_released_with_work(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + + work = source.write(local=[local.region()], remote=[peer.region("buffer")]) + work.wait() + work.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_submit_reservation_is_released_on_base_exception(self) -> None: + source = Endpoint("source", backend="test-interrupt") + destination = Endpoint("destination", backend="test-interrupt") + local = source.register(bytearray(4), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + + with self.assertRaises(KeyboardInterrupt): + source.write(plan) + + self.assertFalse(plan._submitting) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_submit_interruption_after_adoption_retains_lifetimes(self) -> None: + source = Endpoint("source", backend="test-adopt-then-interrupt") + destination = Endpoint("destination", backend="test-adopt-then-interrupt") + local = source.register(bytearray(4), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + + recovery_token = object() + work = source.write(plan, recovery_token=recovery_token) + + self.assertFalse(plan._submitting) + self.assertIs(work.recovery_token, recovery_token) + self.assertIn(work, plan._works) + self.assertIn(work, source._works) + with self.assertRaises(BusyError): + plan.close() + work._backend_work._state = WorkState.COMPLETED + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + + work.close() + self.assertIsNone(work.recovery_token) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_submit_without_work_adoption_is_rolled_back(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(4), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + source._backend.submit = lambda *args, **kwargs: None + + with self.assertRaisesRegex(BackendFailureError, "without adopting"): + source.write(plan) + + self.assertFalse(plan._submitting) + self.assertFalse(plan._works) + self.assertFalse(source._works) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_submit_epilogue_interruption_is_recoverable(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(4), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + + class InterruptOnceOnDiscard(set): + armed = True + + def discard(self, value) -> None: + super().discard(value) + if self.armed: + self.armed = False + raise KeyboardInterrupt + + plan._submitting = InterruptOnceOnDiscard() + work = source.write(plan) + + self.assertFalse(plan._submitting) + self.assertEqual(source.open_works, (work,)) + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) + work.close() + self.assertEqual(source.open_works, ()) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_recovery_token_covers_every_work_producing_surface(self) -> None: + class IdentityOnlyToken: + def __hash__(self): + raise AssertionError("recovery token must not be hashed") + + def __eq__(self, other): + raise AssertionError("recovery token must not be compared") + + def __bool__(self): + raise AssertionError("recovery token must not be coerced") + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"abcd"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("remote", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=array("i", [0]), + remote_indices=array("i", [0]), + ) + bound = selection.bind() + calls = ( + lambda token: plan.submit_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + ), + lambda token: plan.submit_indices_and_test( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + max_polls=0, + ), + lambda token: plan.submit_indices_and_poll( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + max_polls=0, + ), + lambda token: selection.submit(recovery_token=token), + lambda token: selection.submit_and_test(recovery_token=token, max_polls=0), + lambda token: selection.submit_and_poll(recovery_token=token, max_polls=0), + lambda token: bound.submit(recovery_token=token), + lambda token: source.read( + plan, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + ), + lambda token: source.write( + plan, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + ), + lambda token: source.send_notification( + peer, b"recovery", recovery_token=token + ), + ) + + self.assertEqual(len(calls), 10) + for call in calls: + token = IdentityOnlyToken() + result = call(token) + work = result[0] if isinstance(result, tuple) else result + self.assertIs(work.recovery_token, token) + matches = [ + candidate + for candidate in source.open_works + if candidate.recovery_token is token + ] + self.assertEqual(matches, [work]) + work.wait() + work.close() + self.assertIsNone(work.recovery_token) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_recovery_token_is_strong_and_close_commit_is_retryable(self) -> None: + class Token: + pass + + class InterruptOnceOnDiscard(set): + armed = True + + def discard(self, value) -> None: + super().discard(value) + if self.armed: + self.armed = False + raise KeyboardInterrupt("unlink interrupted") + + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + token = Token() + token_ref = weakref.ref(token) + work = source.write(plan, recovery_token=token) + del token + gc.collect() + self.assertIs(token_ref(), work.recovery_token) + work.wait() + + source._works = InterruptOnceOnDiscard(source._works) + with self.assertRaisesRegex(KeyboardInterrupt, "unlink interrupted"): + work.close() + self.assertTrue(work.closed) + self.assertIs(token_ref(), work.recovery_token) + + work.close() + self.assertIsNone(work.recovery_token) + gc.collect() + self.assertIsNone(token_ref()) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_recovery_token_clears_on_every_preadoption_rollback(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"abcd"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("remote", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + ) + bound = plan.bind( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + ) + captured = [] + + def reject_adoption(*args, **kwargs) -> None: + captured.append(kwargs["adopt_work"].__self__) + + cases = ( + ( + "submit", + lambda token: source.write( + plan, + local_indices=[0], + remote_indices=[0], + recovery_token=token, + ), + ), + ("submit", lambda token: bound.submit(recovery_token=token)), + ( + "submit_prevalidated", + lambda token: selection.submit(recovery_token=token), + ), + ( + "send_notification", + lambda token: source.send_notification( + peer, b"rollback", recovery_token=token + ), + ), + ) + for backend_method, submit in cases: + token = object() + with ( + mock.patch.object( + source._backend, + backend_method, + side_effect=reject_adoption, + ), + self.assertRaisesRegex(BackendFailureError, "without adopting"), + ): + submit(token) + self.assertIsNone(captured[-1].recovery_token) + self.assertEqual(source.open_works, ()) + + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_one_shot_validation_failure_does_not_leak_plan(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + + with self.assertRaises(InvalidRegionError): + source.write( + local=[local.region()], + remote=[peer.region("buffer")], + local_indices=[0], + remote_indices=[0], + ) + self.assertFalse(source._plans) + + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_one_shot_plan_release_can_be_retried(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + work = source.write(local=[local.region()], remote=[peer.region("buffer")]) + original_release = source._backend.release_plan + fail_once = True + + def release_plan(plan) -> None: + nonlocal fail_once + if fail_once: + fail_once = False + raise RuntimeError("injected release failure") + original_release(plan) + + source._backend.release_plan = release_plan + with self.assertRaises(BackendFailureError): + work.close() + self.assertFalse(work.closed) + self.assertTrue(source._plans) + work.close() + self.assertFalse(source._plans) + + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_failed_one_shot_cleanup_is_retried_by_child_close(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + original_release = source._backend.release_plan + fail_once = True + + def release_plan(plan) -> None: + nonlocal fail_once + if fail_once: + fail_once = False + raise RuntimeError("injected cleanup failure") + original_release(plan) + + source._backend.release_plan = release_plan + with self.assertRaises(InvalidRegionError): + source.write( + local=[local.region()], + remote=[peer.region("buffer")], + local_indices=[0], + remote_indices=[0], + ) + self.assertTrue(source._orphaned_plans) + peer.close() + self.assertFalse(source._orphaned_plans) + + local.close() + remote.close() + source.close() + destination.close() + + def test_rejects_invalid_regions_and_implicit_copies(self) -> None: + endpoint = Endpoint("endpoint") + registration = endpoint.register(bytearray(8), name="buffer") + with self.assertRaises(InvalidRegionError): + registration.region(7, 2) + with self.assertRaises(InvalidRegionError): + registration.region(0, 1, stride=2) + with self.assertRaises(UnsupportedError): + endpoint.register(RawSpan(1, 8, "cpu", "host", bytearray(8)), name="raw") + + tensor = torch.empty((2, 2)).t() + with self.assertRaises(InvalidRegionError): + endpoint.register(tensor, name="noncontiguous") + with self.assertRaises(InvalidRegionError): + endpoint.register(torch.empty(0), name="empty") + + registration.close() + endpoint.close() + + def test_resized_storage_is_rejected_before_use(self) -> None: + endpoint = Endpoint("endpoint") + tensor = torch.arange(8, dtype=torch.uint8) + storage = tensor.untyped_storage() + registration = endpoint.register(tensor, name="tensor") + storage.resize_(16) + with self.assertRaises(InvalidRegionError): + registration.region() + with self.assertRaises(InvalidRegionError): + endpoint.export_metadata() + registration.close() + endpoint.close() + + def test_cuda_ordering_modes_are_explicit(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + + self.assertEqual( + source.capabilities.cuda_ordering_modes, + frozenset({CudaOrderingMode.CALLER_READY}), + ) + self.assertTrue(source.capabilities.attached_notifications) + self.assertTrue(source.capabilities.standalone_notifications) + self.assertEqual( + source.capabilities.standalone_notification_completion, + NotificationCompletion.REMOTE_DELIVERED, + ) + with self.assertRaises(ValueError): + source.write(plan, stream=object()) + with self.assertRaises(ValueError): + source.write(plan, cuda_ordering="host_wait") + with self.assertRaises(UnsupportedError): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=object(), + ) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_host_wait_is_core_owned_and_runs_outside_backend_lock(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + source_buffer = bytearray(b"old!") + destination_buffer = bytearray(4) + local_registration = source.register(source_buffer, name="buffer") + remote_registration = destination.register(destination_buffer, name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local_registration.region() + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + local_region._device = "cuda:0" + source._capabilities = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + strided_regions=True, + manual_progress=True, + background_progress=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + standalone_notification_completion=( + NotificationCompletion.REMOTE_DELIVERED + ), + ) + + progressed = threading.Event() + test_case = self + + class FakeEvent: + device = torch.device("cuda:0") + + def synchronize(self) -> None: + progress_thread = threading.Thread( + target=lambda: (source.progress(), progressed.set()) + ) + progress_thread.start() + if not progressed.wait(1): + test_case.fail("CUDA wait ran while the provider lock was held") + progress_thread.join(1) + source_buffer[:] = b"new!" + + with ( + mock.patch.object(torch.cuda, "Event", FakeEvent), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=False + ), + ): + work = source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(), + ) + self.assertEqual(destination_buffer, b"new!") + + work.close() + plan.close() + peer.close() + local_registration.close() + remote_registration.close() + source.close() + destination.close() + + def test_host_wait_stream_records_once_and_wait_failure_rolls_back(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local_registration = source.register(bytearray(b"data"), name="buffer") + remote_registration = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local_registration.region() + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + local_region._device = "cuda:1" + source._capabilities = Capabilities( + read=True, + write=True, + manual_progress=True, + background_progress=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + ) + actions: list[str] = [] + + class FakeEvent: + device = torch.device("cuda:1") + + def synchronize(self) -> None: + actions.append("wait") + + class FakeStream: + device = torch.device("cuda:1") + + def record_event(self) -> FakeEvent: + actions.append("record") + return FakeEvent() + + class FakeStreamContext: + def __enter__(self) -> None: + actions.append("capture-check-enter") + + def __exit__(self, *args: object) -> None: + actions.append("capture-check-exit") + + with ( + mock.patch.object(torch, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Event", FakeEvent), + mock.patch.object(torch.cuda, "stream", return_value=FakeStreamContext()), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=False + ), + ): + work = source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + stream=FakeStream(), + ) + self.assertEqual( + actions, + ["capture-check-enter", "capture-check-exit", "record", "wait"], + ) + work.close() + + actions.clear() + with ( + mock.patch.object(torch, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Stream", FakeStream), + mock.patch.object(torch.cuda, "Event", FakeEvent), + mock.patch.object(torch.cuda, "stream", return_value=FakeStreamContext()), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=True + ), + self.assertRaisesRegex(UnsupportedError, "Graph capture"), + ): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + stream=FakeStream(), + ) + self.assertEqual(actions, ["capture-check-enter", "capture-check-exit"]) + self.assertFalse(plan._submitting) + + class FailingEvent(FakeEvent): + def synchronize(self) -> None: + raise RuntimeError("wait failed") + + with ( + mock.patch.object(torch.cuda, "Event", FailingEvent), + mock.patch.object( + torch.cuda, "is_current_stream_capturing", return_value=False + ), + self.assertRaisesRegex(BackendFailureError, "CUDA readiness"), + ): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FailingEvent(), + ) + self.assertFalse(plan._submitting) + + plan.close() + peer.close() + local_registration.close() + remote_registration.close() + source.close() + destination.close() + + def test_host_wait_rejects_wrong_memory_token_and_skips_event_capture_query( + self, + ) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local_registration = source.register(bytearray(b"data"), name="buffer") + remote_registration = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + local_region = local_registration.region() + plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) + source._capabilities = Capabilities( + read=True, + write=True, + manual_progress=True, + background_progress=True, + cuda_ordering_modes=frozenset( + {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} + ), + ) + + class FakeEvent: + def __init__(self, device: object) -> None: + self.device = device + + def synchronize(self) -> None: + pass + + with mock.patch.object(torch.cuda, "Event", FakeEvent): + with self.assertRaisesRegex(InvalidRegionError, "selected local"): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(torch.device("cuda:0")), + ) + local_region._device = "cuda:0" + with self.assertRaisesRegex(InvalidRegionError, "recorded CUDA"): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(None), + ) + with self.assertRaisesRegex(InvalidRegionError, "does not match"): + source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(torch.device("cuda:1")), + ) + with mock.patch.object( + torch.cuda, + "is_current_stream_capturing", + side_effect=AssertionError("ready-event path queried current stream"), + ) as capture_query: + work = source.write( + plan, + cuda_ordering=CudaOrderingMode.HOST_WAIT, + ready_event=FakeEvent(torch.device("cuda:0")), + ) + capture_query.assert_not_called() + work.close() + self.assertFalse(plan._submitting) + + plan.close() + peer.close() + local_registration.close() + remote_registration.close() + source.close() + destination.close() + + def test_expected_peer_identity_is_checked_before_import(self) -> None: + local = Endpoint("local") + remote = Endpoint("remote", endpoint_id="remote-id") + metadata = bytes(remote.export_metadata()) + + with self.assertRaisesRegex(InvalidMetadataError, "name"): + local.import_peer(metadata, expected_name="other") + self.assertEqual(local._peers, {}) + + peer = local.import_peer( + metadata, + expected_identity=EndpointIdentity("remote-id", remote.incarnation), + expected_name="remote", + ) + peer.close() + local.close() + remote.close() + + def test_work_timeout_and_best_effort_cancel(self) -> None: + source = Endpoint("source", backend="test-pending") + destination = Endpoint("destination", backend="test-pending") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + with self.assertRaises(TransferTimeoutError): + work.wait(0) + with self.assertRaises(ValueError): + work.wait(float("nan")) + self.assertTrue(work.cancel()) + self.assertEqual(work.state, WorkState.CANCELLED) + with self.assertRaises(TransferCancelledError): + work.wait() + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_zero_timeout_bounded_work_fallback_observes_once(self) -> None: + backend_work = _CountingPendingWork() + self.assertEqual( + backend_work.poll_bounded( + max_polls=64, + timeout_ns=0, + timeout_check_interval=32, + ), + WorkState.PENDING, + ) + self.assertEqual(backend_work.state_queries, 1) + + for max_polls, timeout_ns, timeout_check_interval in ( + (1 << 63, None, 32), + (1, 1 << 63, 32), + (1, None, 1 << 63), + ): + with self.assertRaises(ValueError): + backend_work.poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + + source = Endpoint("source", backend="test-counting-pending") + destination = Endpoint("destination", backend="test-counting-pending") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + self.assertFalse( + work.test_bounded( + max_polls=64, + timeout_ns=0, + timeout_check_interval=32, + ) + ) + provider_work = source._backend.last_work + self.assertIsNotNone(provider_work) + assert provider_work is not None + self.assertEqual(provider_work.state_queries, 1) + for arguments in ( + {"max_polls": 1 << 63}, + {"timeout_ns": 1 << 63}, + {"timeout_check_interval": 1 << 63}, + ): + with self.assertRaises(ValueError): + work.test_bounded(**arguments) + + self.assertTrue(work.cancel()) + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_work_failure_diagnostic_publication_retries_after_interrupt( + self, + ) -> None: + source = Endpoint("source", backend="test-interrupting-failure") + destination = Endpoint("destination", backend="test-interrupting-failure") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + provider_work = source._backend.last_work + self.assertIsNotNone(provider_work) + assert provider_work is not None + + with self.assertRaises(KeyboardInterrupt): + _ = work.state + self.assertEqual(work._state, WorkState.PENDING) + self.assertIsNone(work._error) + self.assertEqual(provider_work.error_queries, 1) + + self.assertEqual(work.state, WorkState.FAILED) + self.assertIs(work.error, provider_work.failure) + with self.assertRaises(BackendFailureError) as context: + work.wait() + self.assertIs(context.exception.__cause__, provider_work.failure) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_lifecycle_constructors_cannot_be_forged(self) -> None: + endpoint = Endpoint("endpoint") + with self.assertRaises(TypeError): + Work(endpoint, _PendingWork()) + endpoint.close() + + def test_background_only_backend_closes_without_manual_progress(self) -> None: + endpoint = Endpoint( + "endpoint", + backend="test-background-only", + progress_mode="background", + ) + with self.assertRaises(UnsupportedError): + endpoint.progress() + endpoint.close() + + def test_backend_error_must_be_an_exception(self) -> None: + source = Endpoint("source", backend="test-bad-error") + destination = Endpoint("destination", backend="test-bad-error") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + + self.assertEqual(work.state, WorkState.FAILED) + self.assertIsInstance(work.error, BackendFailureError) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_backend_work_state_must_not_regress(self) -> None: + source = Endpoint("source", backend="test-regressing") + destination = Endpoint("destination", backend="test-regressing") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + + self.assertEqual(work.state, WorkState.RUNNING) + with self.assertRaisesRegex(BackendFailureError, "regressed"): + _ = work.state + self.assertEqual(work._state, WorkState.RUNNING) + self.assertTrue(work.cancel()) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_state_query_failure_does_not_fake_terminal_completion(self) -> None: + source = Endpoint("source", backend="test-transient-state-error") + destination = Endpoint("destination", backend="test-transient-state-error") + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + + with self.assertRaises(BackendFailureError): + work.test() + self.assertEqual(work._state, WorkState.PENDING) + with self.assertRaises(BusyError): + work.close() + self.assertFalse(plan.closed) + self.assertFalse(local.closed) + self.assertTrue(work.cancel()) + + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_endpoint_id_detects_reincarnation(self) -> None: + local = Endpoint("local") + first = Endpoint("remote", endpoint_id="stable-remote") + old_peer = local.import_peer(bytes(first.export_metadata())) + first.close() + second = Endpoint("remote", endpoint_id="stable-remote") + + with self.assertRaises(BusyError): + local.import_peer(bytes(second.export_metadata())) + old_peer.close() + new_peer = local.import_peer(bytes(second.export_metadata())) + self.assertNotEqual(old_peer.incarnation, new_peer.incarnation) + + new_peer.close() + local.close() + second.close() + + def test_reference_selective_metadata_is_a_registration_grant(self) -> None: + remote = Endpoint("remote") + first = remote.register(bytearray(4), name="first") + second = remote.register(bytearray(4), name="second") + full = remote.export_metadata() + subset = remote.export_metadata([first]) + local = Endpoint("local") + peer = local.import_peer(bytes(subset)) + with self.assertRaises(InvalidRegionError): + peer.region("second") + + second_metadata = next( + item for item in full.registrations if item.name == "second" + ) + spliced = EndpointMetadata( + wire_major=subset.wire_major, + wire_minor=subset.wire_minor, + required_features=subset.required_features, + backend=subset.backend, + endpoint_id=subset.endpoint_id, + name=subset.name, + incarnation=subset.incarnation, + backend_payload=subset.backend_payload, + registrations=(second_metadata,), + ) + other_local = Endpoint("other-local") + with self.assertRaises(InvalidMetadataError): + other_local.import_peer(bytes(spliced)) + + peer.close() + local.close() + other_local.close() + first.close() + second.close() + remote.close() + + def test_reference_rejects_renamed_registration_metadata(self) -> None: + remote = Endpoint("remote") + remote_registration = remote.register(bytearray(4), name="original") + metadata = remote.export_metadata() + original = metadata.registrations[0] + renamed = EndpointMetadata( + wire_major=metadata.wire_major, + wire_minor=metadata.wire_minor, + required_features=metadata.required_features, + backend=metadata.backend, + endpoint_id=metadata.endpoint_id, + name=metadata.name, + incarnation=metadata.incarnation, + backend_payload=metadata.backend_payload, + registrations=( + RegistrationMetadata( + registration_id=original.registration_id, + name="forged", + nbytes=original.nbytes, + device=original.device, + memory_type=original.memory_type, + backend_descriptor=original.backend_descriptor, + ), + ), + ) + local = Endpoint("local") + local_registration = local.register(bytearray(4), name="local") + peer = local.import_peer(renamed) + + with self.assertRaises(InvalidMetadataError): + local.prepare( + local=[local_registration.region()], + remote=[peer.region("forged")], + ) + + peer.close() + local_registration.close() + local.close() + remote_registration.close() + remote.close() + + def test_single_thread_mode_rejects_other_thread(self) -> None: + endpoint = Endpoint("endpoint", thread_mode="single") + errors: list[BaseException] = [] + + def use_endpoint() -> None: + try: + endpoint.progress() + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=use_endpoint) + thread.start() + thread.join() + self.assertEqual(len(errors), 1) + self.assertIsInstance(errors[0], UnsupportedError) + endpoint.close() + + def test_caller_serialized_mode_allows_sequential_thread_migration(self) -> None: + source = Endpoint("source", thread_mode="caller_serialized") + destination = Endpoint("destination", thread_mode=ThreadMode.CALLER_SERIALIZED) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + work = source.write(plan) + + self.assertIs(source._state_lock, _api._NOOP_LOCK) + self.assertIsNone(work._lock) + self.assertIsNone(work._owner_thread) + failures: list[BaseException] = [] + external_lock = threading.Lock() + + def observe_and_close() -> None: + try: + with external_lock: + self.assertEqual(work.state, WorkState.COMPLETED) + work.close() + source.progress() + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=observe_and_close) + thread.start() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(work.closed) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_single_indexed_success_never_enters_core_noop_lock(self) -> None: + source = Endpoint("source", thread_mode=ThreadMode.SINGLE) + destination = Endpoint("destination") + local = source.register(bytearray(b"abcdefgh"), name="buffer") + remote = destination.register(bytearray(8), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 4, stride=4, count=2)], + remote=[peer.region("buffer", 0, 4, stride=4, count=2)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[1], + remote_indices=[0], + ) + + with mock.patch.object( + _api._NoopLock, + "__enter__", + side_effect=AssertionError("successful indexed path entered _NoopLock"), + ): + work = selection.submit() + self.assertEqual(work.state, WorkState.COMPLETED) + work.close() + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_caller_serialized_exact_slots_migrate_without_core_locks(self) -> None: + source = Endpoint( + "source", + backend="test-request-slot-notification", + thread_mode="caller_serialized", + ) + destination = Endpoint( + "destination", + backend="test-request-slot-notification", + thread_mode="caller_serialized", + ) + local = source.register(bytearray(b"data"), name="buffer") + remote = destination.register(bytearray(4), name="buffer") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + bound = plan.bind(TransferOp.WRITE) + fixed = bound.create_slot() + notification = bound.create_slot(request_local_notifications=True) + + self.assertIsNone(fixed._lock) + self.assertIsNone(fixed._owner_thread) + self.assertIs(fixed._fast_execute, fixed._backend_start_and_poll) + self.assertIsNone(notification._lock) + self.assertIsNone(notification._owner_thread) + self.assertIs( + notification._notification_fast_execute, + notification._backend_start_and_poll_with_notification, + ) + + external_lock = threading.Lock() + + def on_new_thread(function): + results = [] + failures: list[BaseException] = [] + + def invoke() -> None: + try: + with external_lock: + results.append(function()) + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=invoke) + thread.start() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(len(results), 1) + return results[0] + + with mock.patch.object( + _api, + "_GET_IDENT", + side_effect=AssertionError( + "caller-serialized slot queried an unowned thread identity" + ), + ): + self.assertEqual(on_new_thread(fixed.execute), 1) + fixed_backend = source._backend.request_slots[0] + fixed_backend.next_state = WorkState.RUNNING + self.assertFalse(on_new_thread(fixed.start_and_test)) + fixed_backend.poll_states = [WorkState.COMPLETED] + self.assertTrue(on_new_thread(fixed.test_and_recycle)) + + self.assertEqual( + on_new_thread(lambda: notification.execute(notification=b"first")), + 1, + ) + notification_backend = source._backend.request_slots[1] + notification_backend.next_state = WorkState.RUNNING + self.assertFalse( + on_new_thread(lambda: notification.start_and_test(notification=None)) + ) + notification_backend.poll_states = [WorkState.COMPLETED] + self.assertTrue(on_new_thread(notification.test_and_recycle)) + + self.assertIs(fixed._fast_execute, fixed._backend_start_and_poll) + self.assertIs( + notification._notification_fast_execute, + notification._backend_start_and_poll_with_notification, + ) + notification.close() + fixed.close() + bound.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_metadata_rejects_unknown_fields(self) -> None: + endpoint = Endpoint("endpoint") + encoded = bytes(endpoint.export_metadata()) + corrupted = encoded[:-1] + b',"unknown":1}' + with self.assertRaises(InvalidMetadataError): + EndpointMetadata.from_bytes(corrupted) + overlong_integer = encoded.replace(b'"major":1', b'"major":' + b"9" * 4301, 1) + with self.assertRaises(InvalidMetadataError): + EndpointMetadata.from_bytes(overlong_integer) + with self.assertRaises(InvalidMetadataError): + EndpointMetadata.from_bytes(b"[" * 2000 + b"0" + b"]" * 2000) + endpoint.close() + + def test_metadata_wire_minor_extensions_and_bounds(self) -> None: + endpoint = Endpoint("endpoint") + current = endpoint.export_metadata() + future_optional = EndpointMetadata( + wire_major=current.wire_major, + wire_minor=current.wire_minor + 1, + required_features=(), + backend=current.backend, + endpoint_id=current.endpoint_id, + name=current.name, + incarnation=current.incarnation, + backend_payload=current.backend_payload, + registrations=current.registrations, + extensions={"future.optional": b"opaque"}, + ) + decoded = EndpointMetadata.from_bytes(bytes(future_optional)) + self.assertEqual(decoded.wire_minor, current.wire_minor + 1) + self.assertEqual(decoded.extensions, (("future.optional", b"opaque"),)) + + with self.assertRaisesRegex(InvalidMetadataError, "required"): + EndpointMetadata( + wire_major=current.wire_major, + wire_minor=current.wire_minor + 1, + required_features=("future.required",), + backend=current.backend, + endpoint_id=current.endpoint_id, + name=current.name, + incarnation=current.incarnation, + backend_payload=current.backend_payload, + registrations=current.registrations, + ) + with self.assertRaisesRegex(InvalidMetadataError, "size limit"): + EndpointMetadata.from_bytes(b" " * (EndpointMetadata.MAX_ENCODED_BYTES + 1)) + + multidimensional = memoryview(bytearray(b" ")).cast("B", (1, 4)) + self.assertEqual(len(multidimensional), 1) + self.assertEqual(multidimensional.nbytes, 4) + with ( + mock.patch.object(EndpointMetadata, "MAX_ENCODED_BYTES", 3), + self.assertRaisesRegex(InvalidMetadataError, "size limit"), + ): + EndpointMetadata.from_bytes(multidimensional) + + registration = RegistrationMetadata( + registration_id="r", + name="n", + nbytes=1, + device="cuda", + memory_type="m", + backend_descriptor=b"", + ) + with ( + mock.patch.object(EndpointMetadata, "MAX_STRING_BYTES", 3), + self.assertRaisesRegex(InvalidMetadataError, "registration device"), + ): + EndpointMetadata( + wire_major=current.wire_major, + wire_minor=current.wire_minor, + required_features=(), + backend="ref", + endpoint_id="id", + name="n", + incarnation="i", + backend_payload=b"", + registrations=(registration,), + ) + + with ( + mock.patch.object(EndpointMetadata, "MAX_EXTENSION_BYTES", 1), + self.assertRaisesRegex(InvalidMetadataError, "extension.*size limit"), + ): + EndpointMetadata( + wire_major=current.wire_major, + wire_minor=current.wire_minor, + required_features=(), + backend=current.backend, + endpoint_id=current.endpoint_id, + name=current.name, + incarnation=current.incarnation, + backend_payload=current.backend_payload, + registrations=current.registrations, + extensions={"optional": b"xx"}, + ) + + with self.assertRaisesRegex(InvalidMetadataError, "valid Unicode"): + EndpointMetadata( + wire_major=current.wire_major, + wire_minor=current.wire_minor, + required_features=(), + backend=current.backend, + endpoint_id=current.endpoint_id, + name="invalid-\ud800", + incarnation=current.incarnation, + backend_payload=current.backend_payload, + registrations=current.registrations, + ) + endpoint.close() + + def test_core_has_no_exception_unsafe_manual_lock_acquisition(self) -> None: + tree = ast.parse(Path(_api.__file__).read_text(encoding="utf-8")) + acquisitions = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "acquire" + ] + self.assertEqual(acquisitions, []) + + def test_first_publication_interruptions_roll_back_every_resource_family( + self, + ) -> None: + class InterruptAfterAdd(set): + def __init__(self, values=()) -> None: + super().__init__(values) + self.armed = True + self.last = None + + def add(self, value) -> None: + super().add(value) + self.last = value + if self.armed: + self.armed = False + raise KeyboardInterrupt("post-add publication cut") + + class InterruptAfterSet(dict): + def __init__(self, values=()) -> None: + super().__init__(values) + self.armed = True + self.last = None + + def __setitem__(self, key, value) -> None: + super().__setitem__(key, value) + self.last = value + if self.armed: + self.armed = False + raise KeyboardInterrupt("post-set publication cut") + + endpoint = Endpoint("registration-publication") + for attribute, name in ( + ("_registrations", "primary"), + ("_registration_names", "name-index"), + ): + collection = InterruptAfterSet(getattr(endpoint, attribute)) + setattr(endpoint, attribute, collection) + with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): + endpoint.register(bytearray(4), name=name) + self.assertTrue(collection.last.closed) + self.assertEqual(endpoint.open_registrations, ()) + self.assertFalse(endpoint._registration_names) + endpoint.close() + + source = Endpoint("source-publication") + destination = Endpoint("destination-publication") + local = source.register(bytearray(b"abcd"), name="local") + remote = destination.register(bytearray(4), name="remote") + metadata = bytes(destination.export_metadata()) + + peer_publication = InterruptAfterSet(source._peers) + source._peers = peer_publication + with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): + source.import_peer(metadata) + self.assertTrue(peer_publication.last.closed) + self.assertEqual(source.open_peers, ()) + peer = source.import_peer(metadata) + + plan_publication = InterruptAfterAdd(source._plans) + source._plans = plan_publication + with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): + source.prepare(local=[local.region()], remote=[peer.region("remote")]) + self.assertTrue(plan_publication.last.closed) + self.assertEqual(source.open_plans, ()) + + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("remote", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + ) + for public_first in (True, False): + if public_first: + publication = InterruptAfterAdd(source._bound_transfers) + source._bound_transfers = publication + else: + publication = InterruptAfterAdd(plan._bindings) + plan._bindings = publication + with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): + selection.bind() + self.assertTrue(publication.last.closed) + self.assertEqual(source.open_bound_transfers, ()) + self.assertFalse(plan._bindings) + + work_publication = InterruptAfterAdd(source._works) + source._works = work_publication + token = object() + with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): + source.send_notification(peer, b"not-sent", recovery_token=token) + self.assertTrue(work_publication.last.closed) + self.assertFalse(work_publication.last._submitting) + self.assertIsNone(work_publication.last.recovery_token) + self.assertEqual(source.open_works, ()) + + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_published_work_is_observation_gated_until_submit_reconciles(self) -> None: + for indexed in (False, True): + with self.subTest(indexed=indexed): + source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) + destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) + local = source.register(bytearray(b"abcd"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("remote", 0, 1, stride=1, count=4)], + indexed=indexed, + ) + backend_work = _CountingPendingWork() + entered = threading.Event() + release = threading.Event() + + def blocked_submit(*args, **kwargs) -> None: + kwargs["adopt_work"](backend_work) + entered.set() + release.wait(2) + + if indexed: + source._backend.submit_prevalidated = blocked_submit + selection = plan.select_indices( + TransferOp.WRITE, + local_indices=[0], + remote_indices=[0], + ) + submit = selection.submit + else: + source._backend.submit = blocked_submit + submit = lambda *, recovery_token: source.write( + plan, recovery_token=recovery_token + ) + token = object() + returned: list[Work] = [] + submit_thread = threading.Thread( + target=lambda: returned.append(submit(recovery_token=token)) + ) + submit_thread.start() + self.assertTrue(entered.wait(1)) + try: + (work,) = source.open_works + self.assertIs(work.recovery_token, token) + observations = ( + lambda: work.state, + lambda: work.error, + work.test, + lambda: work.poll_bounded(max_polls=1, timeout_ns=None), + work.cancel, + lambda: work.completion_event, + work.close, + ) + for observation in observations: + with self.assertRaisesRegex( + BusyError, "submission is in progress" + ): + observation() + self.assertEqual(backend_work.state_queries, 0) + with self.assertRaisesRegex(BusyError, "submissions in progress"): + plan.close() + finally: + release.set() + submit_thread.join(2) + self.assertFalse(submit_thread.is_alive()) + self.assertEqual(returned, [work]) + self.assertFalse(work._submitting) + backend_work._state = WorkState.COMPLETED + self.assertEqual(work.state, WorkState.COMPLETED) + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_indexed_publication_is_owned_before_lock_exit_and_every_graph_edge( + self, + ) -> None: + class InterruptAfterAdd(set): + def __init__(self, values=()) -> None: + super().__init__(values) + self.armed = True + self.last = None + + def add(self, value) -> None: + super().add(value) + self.last = value + if self.armed: + self.armed = False + raise KeyboardInterrupt("indexed graph publication cut") + + class InterruptAfterExit: + def __init__(self, lock) -> None: + self.lock = lock + self.armed = True + + def __enter__(self): + self.lock.__enter__() + return self + + def __exit__(self, exc_type, exc_value, traceback): + result = self.lock.__exit__(exc_type, exc_value, traceback) + if self.armed: + self.armed = False + raise KeyboardInterrupt("indexed state-lock exit cut") + return result + + source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) + destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) + local = source.register(bytearray(b"abcd"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare( + local=[local.region(0, 1, stride=1, count=4)], + remote=[peer.region("remote", 0, 1, stride=1, count=4)], + indexed=True, + ) + selection = plan.select_indices( + TransferOp.WRITE, local_indices=[0], remote_indices=[0] + ) + backend_entries = 0 + + def unexpected_backend_entry(*args, **kwargs) -> None: + nonlocal backend_entries + backend_entries += 1 + raise AssertionError("publication interruption reached provider") + + source._backend.submit_prevalidated = unexpected_backend_entry + original_state_lock = source._state_lock + source._state_lock = InterruptAfterExit(original_state_lock) + try: + with self.assertRaisesRegex(KeyboardInterrupt, "state-lock exit cut"): + selection.submit(recovery_token=object()) + finally: + source._state_lock = original_state_lock + self.assertEqual(source.open_works, ()) + self.assertFalse(plan._submitting) + self.assertFalse(plan._works) + + for owner, attribute in ( + (source, "_works"), + (plan, "_submitting"), + (plan, "_works"), + ): + with self.subTest(edge=f"{type(owner).__name__}.{attribute}"): + original = getattr(owner, attribute) + publication = InterruptAfterAdd(original) + setattr(owner, attribute, publication) + try: + with self.assertRaisesRegex( + KeyboardInterrupt, "graph publication cut" + ): + selection.submit(recovery_token=object()) + self.assertIsNotNone(publication.last) + self.assertTrue(publication.last.closed) + self.assertFalse(publication.last._submitting) + self.assertIsNone(publication.last.recovery_token) + finally: + setattr(owner, attribute, set(publication)) + self.assertEqual(source.open_works, ()) + self.assertFalse(plan._submitting) + self.assertFalse(plan._works) + + self.assertEqual(backend_entries, 0) + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_endpoint_post_adoption_initialization_cleanup_is_complete(self) -> None: + for cut in ("capabilities", "dispatch", "endpoint-info"): + with self.subTest(cut=cut): + instances = [] + + class InitializationBackend(_ReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.close_calls = 0 + + @property + def capabilities(self): + if cut == "capabilities": + raise KeyboardInterrupt("capabilities cut") + return super().capabilities + + def __getattribute__(self, name): + if name == "send_notification_sync": + current_cut = object.__getattribute__(self, "_test_cut") + if current_cut == "dispatch": + raise KeyboardInterrupt("dispatch cut") + return super().__getattribute__(name) + + def close(self) -> None: + self.close_calls += 1 + return super().close() + + def factory(*args, **kwargs): + implementation = InitializationBackend(*args, **kwargs) + implementation._test_cut = cut + instances.append(implementation) + return implementation + + backend_name = f"test-endpoint-init-cleanup-{cut}" + register_backend(backend_name, factory, replace=True) + if cut == "endpoint-info": + with ( + mock.patch.object( + _api, + "EndpointInfo", + side_effect=KeyboardInterrupt("EndpointInfo cut"), + ), + self.assertRaisesRegex(KeyboardInterrupt, "EndpointInfo cut"), + ): + Endpoint("interrupted", backend=backend_name) + else: + with self.assertRaisesRegex(KeyboardInterrupt, f"{cut} cut"): + Endpoint("interrupted", backend=backend_name) + self.assertEqual(len(instances), 1) + self.assertEqual(instances[0].close_calls, 1) + self.assertTrue(instances[0]._closed) + + def test_endpoint_failed_initialization_exposes_retryable_cleanup_owner( + self, + ) -> None: + marker = object() + instances = [] + + class HeldInitializationBackend(_ReferenceBackend): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + super().close() + return marker if self.close_calls == 1 else None + + def factory(*args, **kwargs): + implementation = HeldInitializationBackend(*args, **kwargs) + instances.append(implementation) + return implementation + + backend_name = "test-endpoint-init-held-cleanup" + register_backend(backend_name, factory, replace=True) + with ( + mock.patch.object( + _api, + "EndpointInfo", + side_effect=KeyboardInterrupt("EndpointInfo cut"), + ), + self.assertRaisesRegex(BackendFailureError, "recovery_endpoint") as context, + ): + Endpoint("interrupted", backend=backend_name) + + recovery_endpoint = context.exception.recovery_endpoint + self.assertIs(recovery_endpoint._backend, instances[0]) + self.assertFalse(recovery_endpoint.closed) + self.assertTrue(recovery_endpoint._creation_failed) + self.assertEqual(instances[0].close_calls, 1) + with self.assertRaisesRegex(BackendFailureError, "creation cleanup only"): + recovery_endpoint.register(bytearray(1)) + recovery_endpoint.close() + self.assertTrue(recovery_endpoint.closed) + self.assertFalse(recovery_endpoint._creation_failed) + self.assertEqual(instances[0].close_calls, 2) + + def test_submit_success_epilogues_keep_noop_lock_fast_branches(self) -> None: + source = Path(_api.__file__).read_text(encoding="utf-8") + generic = source[source.index(" def _submit(\n") :] + generic = generic[: generic.index(" def _validate_regions(\n")] + bound = source[source.index(" def _submit_bound(\n") :] + bound = bound[: bound.index(" def send_notification(\n")] + self.assertIn("if state_lock is _NOOP_LOCK:", generic) + self.assertIn("if state_lock is _NOOP_LOCK:", bound) + poll = source[source.index(" def poll_notification_batches(\n") :] + poll = poll[: poll.index(" def progress(")] + conditional = poll.index("if max_items is not None:") + payload_count = poll.index( + "payload_count = sum(len(batch.payloads) for batch in batches)" + ) + self.assertLess(conditional, payload_count) + + def test_teardown_requires_exact_none_before_core_commit(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + work = source.write(plan) + self.assertEqual(work.state, WorkState.COMPLETED) + + marker = object() + + def expect_retryable(resource, owner, attribute, original, *args) -> None: + def close_then_return_marker(*call_args, **call_kwargs): + original(*call_args, **call_kwargs) + return marker + + setattr(owner, attribute, close_then_return_marker) + with self.assertRaisesRegex(BackendFailureError, "must return None"): + resource.close() + self.assertFalse(resource.closed) + setattr(owner, attribute, original) + resource.close() + + backend_work = work._backend_work + expect_retryable(work, backend_work, "close", backend_work.close) + expect_retryable( + plan, + source._backend, + "release_plan", + source._backend.release_plan, + plan._handle, + ) + expect_retryable( + peer, + source._backend, + "release_peer", + source._backend.release_peer, + peer._handle, + ) + expect_retryable( + local, + source._backend, + "deregister", + source._backend.deregister, + local._handle, + ) + remote.close() + original_close = source._backend.close + expect_retryable(source, source._backend, "close", original_close) + destination.close() + + def test_work_completion_event_is_capability_gated(self) -> None: + source = Endpoint("source") + destination = Endpoint("destination") + local = source.register(bytearray(b"data"), name="local") + remote = destination.register(bytearray(4), name="remote") + peer = source.import_peer(bytes(destination.export_metadata())) + plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) + work = source.write(plan) + with mock.patch.object( + type(work._backend_work), + "completion_event", + new_callable=mock.PropertyMock, + side_effect=AssertionError("disabled capability entered provider"), + ): + self.assertIsNone(work.completion_event) + work.close() + plan.close() + peer.close() + local.close() + remote.close() + source.close() + destination.close() + + def test_notification_batch_exact_type_fast_path_skips_payload_rescan(self) -> None: + endpoint = Endpoint("receiver") + payloads = (b"one", b"two") + batch = BackendNotificationBatch(payloads) + endpoint._backend.poll_notification_batches = lambda max_items: [batch] + with mock.patch.object( + _types, + "_immutable_notification_payloads", + side_effect=AssertionError("validated backend tuple was rescanned"), + ): + (public,) = endpoint.poll_notification_batches() + self.assertIs(public.payloads, payloads) + + class BackendBatchSubclass(BackendNotificationBatch): + pass + + endpoint._backend.poll_notification_batches = lambda max_items: [ + BackendBatchSubclass((b"subclass",)) + ] + with self.assertRaisesRegex(BackendFailureError, "must be a list"): + endpoint.poll_notification_batches() + endpoint.close() + + def test_provider_objects_redact_owners_and_capability_descriptors(self) -> None: + secret = b"provider-capability-secret" + owner = {"secret": secret} + registered = RegisteredMemory("id", "name", owner, 1, 1, "cpu", "host") + imported = ImportedRegistration( + owner, + "id", + "name", + 1, + "cpu", + "host", + secret, + ) + raw = ImportedRawMemory(owner, 1, 1, "cpu", "host") + region = BackendRegion( + owner=owner, + offset=0, + nbytes=1, + device="cpu", + memory_type="host", + backend_descriptor=secret, + ) + + for value in (registered, imported, raw, region): + self.assertNotIn("provider-capability-secret", repr(value)) + + +if __name__ == "__main__": + run_tests() diff --git a/test/dynamo/test_trace_rules.py b/test/dynamo/test_trace_rules.py index d7c27c6efd291..1d4d7206e1113 100644 --- a/test/dynamo/test_trace_rules.py +++ b/test/dynamo/test_trace_rules.py @@ -191,6 +191,7 @@ def _is_allowed_module_prefix(obj): "torch._C._distributed_autograd", "torch._C._distributed_c10d", "torch._C._distributed_rpc", + "torch._C._distributed_transfer", "torch._C._functorch", "torch._C._monitor", "torch._C._nvtx", diff --git a/torch/_C/_distributed_transfer.pyi b/torch/_C/_distributed_transfer.pyi new file mode 100644 index 0000000000000..b279648311126 --- /dev/null +++ b/torch/_C/_distributed_transfer.pyi @@ -0,0 +1,7 @@ +# This module is defined in torch/csrc/distributed/transfer/python_bindings.cpp + +def _snapshot_and_analyze_int32_indices( + local_indices: object, + remote_indices: object, + destination_is_remote: bool, +) -> tuple[bytes, int, int, int, int, int, int, bool]: ... diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp index 200dcebdf572d..bd384cacc35e6 100644 --- a/torch/csrc/Module.cpp +++ b/torch/csrc/Module.cpp @@ -73,6 +73,7 @@ #include #include #include +#include #include #include #include @@ -2669,6 +2670,7 @@ PyObject* initModule() { #endif torch::distributed::initPlacementBindings(module); + torch::distributed::transfer::initPythonBindings(module); torch::initTensorIteratorBindings(module); diff --git a/torch/csrc/distributed/transfer/python_bindings.cpp b/torch/csrc/distributed/transfer/python_bindings.cpp new file mode 100644 index 0000000000000..4969acd80cf71 --- /dev/null +++ b/torch/csrc/distributed/transfer/python_bindings.cpp @@ -0,0 +1,432 @@ +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace torch::distributed::transfer { +namespace { + +constexpr std::size_t kTinyDuplicateThreshold = 32; +constexpr std::uint64_t kMaximumBitmapValueCount = 1ULL << 24; +constexpr std::uint64_t kMaximumBitmapValuesPerIndex = 16; + +class BufferView final { + public: + BufferView(PyObject* object, const char* name) : name_(name) { + if (PyObject_GetBuffer(object, &view_, PyBUF_FORMAT | PyBUF_STRIDES) != 0) { + throw py::error_already_set(); + } + acquired_ = true; + try { + validate(); + } catch (...) { + release(); + throw; + } + } + + BufferView(const BufferView&) = delete; + BufferView& operator=(const BufferView&) = delete; + + ~BufferView() { + release(); + } + + const void* data() const { + return view_.buf; + } + + Py_ssize_t nbytes() const { + return view_.len; + } + + std::size_t count() const { + return static_cast(view_.len) / sizeof(std::int32_t); + } + + private: + static char nativeEndianPrefix() { + const std::uint16_t value = 1; + std::uint8_t first_byte = 0; + std::memcpy(&first_byte, &value, sizeof(first_byte)); + return first_byte == 1 ? '<' : '>'; + } + + static bool isNativeInt32Format(const char* format) { + if (format == nullptr) { + return false; + } + if (std::strcmp(format, "i") == 0 || std::strcmp(format, "@i") == 0 || + std::strcmp(format, "=i") == 0) { + return true; + } + return format[0] == nativeEndianPrefix() && + std::strcmp(format + 1, "i") == 0; + } + + void validate() const { + if (view_.ndim != 1 || view_.itemsize != sizeof(std::int32_t) || + view_.len < 0 || view_.len % sizeof(std::int32_t) != 0 || + !isNativeInt32Format(view_.format)) { + throw py::value_error( + std::string(name_) + + " must be a one-dimensional C-contiguous native-int32 buffer"); + } + const auto count = view_.len / sizeof(std::int32_t); + if (view_.buf == nullptr || view_.shape == nullptr || + view_.strides == nullptr || view_.shape[0] != count || + (view_.suboffsets != nullptr && view_.suboffsets[0] >= 0)) { + throw py::value_error( + std::string(name_) + " has an inconsistent PEP 3118 shape"); + } + if (PyBuffer_IsContiguous(&view_, 'C') == 0) { + throw py::value_error( + std::string(name_) + + " must be a one-dimensional C-contiguous native-int32 buffer"); + } + if (count == 0) { + throw py::value_error(std::string(name_) + " must be non-empty"); + } + if (static_cast(view_.len) > + std::numeric_limits::max()) { + throw std::overflow_error( + std::string(name_) + " is too large for this platform"); + } + } + + void release() { + if (acquired_) { + PyBuffer_Release(&view_); + acquired_ = false; + } + } + + const char* name_; + Py_buffer view_{}; + bool acquired_{false}; +}; + +struct IndexAnalysis final { + std::int32_t local_minimum; + std::int32_t local_maximum; + std::int32_t remote_minimum; + std::int32_t remote_maximum; + bool destination_has_duplicate; +}; + +struct BoundsAnalysis final { + std::int32_t minimum; + std::int32_t maximum; +}; + +struct DestinationAnalysis final { + BoundsAnalysis bounds; + bool strictly_monotonic; +}; + +std::int32_t loadInt32(const char* data, std::size_t index) { + std::int32_t value = 0; + std::memcpy( + &value, data + index * sizeof(std::int32_t), sizeof(std::int32_t)); + return value; +} + +std::uint64_t mixKey(std::uint64_t value) { + value += 0x9e3779b97f4a7c15ULL; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +bool hasTinyDuplicate(const char* data, std::size_t count) { + for (std::size_t index = 1; index < count; ++index) { + const auto value = loadInt32(data, index); + for (std::size_t prior = 0; prior < index; ++prior) { + if (loadInt32(data, prior) == value) { + return true; + } + } + } + return false; +} + +bool hasBitmapDuplicate( + const char* data, + std::size_t count, + std::int32_t minimum, + std::uint64_t value_count) { + std::vector bitmap( + static_cast((value_count + 63) / 64), 0); + for (std::size_t index = 0; index < count; ++index) { + const auto offset = static_cast( + static_cast(loadInt32(data, index)) - minimum); + const auto word = static_cast(offset / 64); + const auto mask = std::uint64_t{1} << (offset % 64); + if ((bitmap[word] & mask) != 0) { + return true; + } + bitmap[word] |= mask; + } + return false; +} + +bool hasFlatHashDuplicate(const char* data, std::size_t count) { + if (count > std::numeric_limits::max() / 2) { + throw std::overflow_error("int32 index count exceeds hash-table capacity"); + } + const auto requested_capacity = count * 2; + std::size_t capacity = 1; + while (capacity < requested_capacity) { + if (capacity > std::numeric_limits::max() / 2) { + throw std::overflow_error("int32 index hash-table capacity overflow"); + } + capacity *= 2; + } + std::vector table(capacity, 0); + const auto mask = capacity - 1; + for (std::size_t index = 0; index < count; ++index) { + const auto key = static_cast( + static_cast(loadInt32(data, index))) + + 1; + auto slot = static_cast(mixKey(key)) & mask; + while (table[slot] != 0) { + if (table[slot] == key) { + return true; + } + slot = (slot + 1) & mask; + } + table[slot] = key; + } + return false; +} + +bool hasDuplicate( + const char* data, + std::size_t count, + std::int32_t minimum, + std::int32_t maximum, + bool strictly_monotonic) { + if (count < 2 || strictly_monotonic) { + return false; + } + const auto value_count = + static_cast(static_cast(maximum) - minimum) + + 1; + if (value_count < count) { + return true; + } + if (count <= kTinyDuplicateThreshold) { + return hasTinyDuplicate(data, count); + } + const auto dense_enough = + count >= (value_count + kMaximumBitmapValuesPerIndex - 1) / + kMaximumBitmapValuesPerIndex; + if (value_count <= kMaximumBitmapValueCount && dense_enough) { + return hasBitmapDuplicate(data, count, minimum, value_count); + } + return hasFlatHashDuplicate(data, count); +} + +BoundsAnalysis analyzeBounds(const char* data, std::size_t count) { + const auto first = loadInt32(data, 0); + BoundsAnalysis result{first, first}; + for (std::size_t index = 1; index < count; ++index) { + const auto value = loadInt32(data, index); + result.minimum = std::min(result.minimum, value); + result.maximum = std::max(result.maximum, value); + } + return result; +} + +DestinationAnalysis finishNonmonotonicDestination( + const char* data, + std::size_t count, + std::size_t suffix_index, + BoundsAnalysis bounds) { + if (suffix_index < count) { + const auto suffix = analyzeBounds( + data + suffix_index * sizeof(std::int32_t), count - suffix_index); + bounds.minimum = std::min(bounds.minimum, suffix.minimum); + bounds.maximum = std::max(bounds.maximum, suffix.maximum); + } + return {bounds, false}; +} + +DestinationAnalysis analyzeDestination(const char* data, std::size_t count) { + // Real framework selections are commonly ordered. Prove their direction + // with a short scalar probe, then let the compiler vectorize the long tail's + // order and bounds reductions. Non-monotonic inputs leave the probe quickly + // and retain the same exact duplicate-analysis fallbacks. + constexpr std::size_t kMonotonicProbeElementCount = 8; + const auto first = loadInt32(data, 0); + DestinationAnalysis result{{first, first}, true}; + if (count == 1) { + return result; + } + + auto previous = loadInt32(data, 1); + if (previous == first) { + return finishNonmonotonicDestination(data, count, 2, result.bounds); + } + + const auto probe_end = std::min(count, kMonotonicProbeElementCount); + std::size_t index = 2; + if (previous > first) { + result.bounds.maximum = previous; + for (; index < probe_end; ++index) { + const auto value = loadInt32(data, index); + if (value <= previous) { + result.bounds.minimum = std::min(result.bounds.minimum, value); + return finishNonmonotonicDestination( + data, count, index + 1, result.bounds); + } + result.bounds.maximum = value; + previous = value; + } + + auto minimum = result.bounds.minimum; + auto maximum = result.bounds.maximum; + bool strictly_monotonic = true; + for (; index < count; ++index) { + const auto value = loadInt32(data, index); + minimum = std::min(minimum, value); + maximum = std::max(maximum, value); + strictly_monotonic &= loadInt32(data, index - 1) < value; + } + return {{minimum, maximum}, strictly_monotonic}; + } + + result.bounds.minimum = previous; + for (; index < probe_end; ++index) { + const auto value = loadInt32(data, index); + if (value >= previous) { + result.bounds.maximum = std::max(result.bounds.maximum, value); + return finishNonmonotonicDestination( + data, count, index + 1, result.bounds); + } + result.bounds.minimum = value; + previous = value; + } + + auto minimum = result.bounds.minimum; + auto maximum = result.bounds.maximum; + bool strictly_monotonic = true; + for (; index < count; ++index) { + const auto value = loadInt32(data, index); + minimum = std::min(minimum, value); + maximum = std::max(maximum, value); + strictly_monotonic &= loadInt32(data, index - 1) > value; + } + return {{minimum, maximum}, strictly_monotonic}; +} + +IndexAnalysis analyzeSnapshot( + const char* data, + std::size_t local_count, + std::size_t remote_count, + bool destination_is_remote) { + const auto* remote = data + local_count * sizeof(std::int32_t); + BoundsAnalysis local_bounds{}; + BoundsAnalysis remote_bounds{}; + DestinationAnalysis destination{}; + if (destination_is_remote) { + local_bounds = analyzeBounds(data, local_count); + destination = analyzeDestination(remote, remote_count); + remote_bounds = destination.bounds; + } else { + destination = analyzeDestination(data, local_count); + local_bounds = destination.bounds; + remote_bounds = analyzeBounds(remote, remote_count); + } + IndexAnalysis result{ + local_bounds.minimum, + local_bounds.maximum, + remote_bounds.minimum, + remote_bounds.maximum, + false}; + const auto* destination_data = destination_is_remote ? remote : data; + const auto destination_count = + destination_is_remote ? remote_count : local_count; + result.destination_has_duplicate = hasDuplicate( + destination_data, + destination_count, + destination.bounds.minimum, + destination.bounds.maximum, + destination.strictly_monotonic); + return result; +} + +py::tuple snapshotAndAnalyzeInt32Indices( + py::object local_indices, + py::object remote_indices, + bool destination_is_remote) { + py::bytes snapshot; + std::size_t local_count = 0; + std::size_t remote_count = 0; + { + BufferView local(local_indices.ptr(), "local_indices"); + BufferView remote(remote_indices.ptr(), "remote_indices"); + if (local.nbytes() > PY_SSIZE_T_MAX - remote.nbytes()) { + throw std::overflow_error("combined int32 index snapshot is too large"); + } + const auto snapshot_nbytes = local.nbytes() + remote.nbytes(); + PyObject* raw_snapshot = + PyBytes_FromStringAndSize(nullptr, snapshot_nbytes); + if (raw_snapshot == nullptr) { + throw py::error_already_set(); + } + snapshot = py::reinterpret_steal(raw_snapshot); + auto* snapshot_data = PyBytes_AS_STRING(snapshot.ptr()); + std::memcpy(snapshot_data, local.data(), local.nbytes()); + std::memcpy(snapshot_data + local.nbytes(), remote.data(), remote.nbytes()); + local_count = local.count(); + remote_count = remote.count(); + } + + const auto* snapshot_data = PyBytes_AS_STRING(snapshot.ptr()); + IndexAnalysis analysis{}; + { + py::gil_scoped_release release; + analysis = analyzeSnapshot( + snapshot_data, local_count, remote_count, destination_is_remote); + } + return py::make_tuple( + std::move(snapshot), + local_count, + remote_count, + analysis.local_minimum, + analysis.local_maximum, + analysis.remote_minimum, + analysis.remote_maximum, + analysis.destination_has_duplicate); +} + +} // namespace + +void initPythonBindings(PyObject* module) { + auto root = py::handle(module).cast(); + auto transfer = root.def_submodule( + "_distributed_transfer", "Private distributed transfer helpers."); + transfer.def( + "_snapshot_and_analyze_int32_indices", + &snapshotAndAnalyzeInt32Indices, + py::arg("local_indices"), + py::arg("remote_indices"), + py::arg("destination_is_remote"), + R"doc( +Create one immutable packed snapshot of two native-int32 buffers and analyze it. + +This is a private implementation helper for ``torch.distributed._transfer``. +)doc"); +} + +} // namespace torch::distributed::transfer diff --git a/torch/csrc/distributed/transfer/python_bindings.h b/torch/csrc/distributed/transfer/python_bindings.h new file mode 100644 index 0000000000000..7518e726c9717 --- /dev/null +++ b/torch/csrc/distributed/transfer/python_bindings.h @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace torch::distributed::transfer { + +void initPythonBindings(PyObject* module); + +} // namespace torch::distributed::transfer diff --git a/torch/distributed/_transfer/README.md b/torch/distributed/_transfer/README.md new file mode 100644 index 0000000000000..472e09c4dda0c --- /dev/null +++ b/torch/distributed/_transfer/README.md @@ -0,0 +1,814 @@ +# Experimental endpoint transfer + +`torch.distributed._transfer` is a private, backend-neutral data-plane API for +moving registered tensor or buffer bytes between independently managed +endpoints. It is intended for inference KV movement and live weight streaming, +where collective process-group semantics are the wrong abstraction. + +NIXL is the first external provider. PyTorch does not import, link, or otherwise +depend on NIXL. The in-process CPU `reference` backend exists for contract tests, +not performance. + +## Scope + +The layer owns: + +- allocation lifetime anchors and byte-range validation; +- versioned endpoint and named-registration metadata; +- peer, plan, work, binding, and persistent-request-slot resource ordering; +- READ, WRITE, indexed transfer, progress, cancellation, and notifications; +- typed errors and explicit backend capabilities; +- an experimental Python provider SPI. + +Frameworks continue to own discovery, routing, leases, cache/model versioning, +layout policy, scheduling, and failure recovery. This package is not a new +control plane and does not provide durable checkpointing. + +## Basic flow + +```python +from torch.distributed import _transfer + +source = _transfer.Endpoint("source", backend="reference") +destination = _transfer.Endpoint("destination", backend="reference") + +source_registration = source.register(source_tensor, name="kv") +destination_registration = destination.register(destination_tensor, name="kv") + +# The application transports these opaque bytes over its control plane. +destination_metadata = bytes(destination.export_metadata()) +destination_peer = source.import_peer(destination_metadata) + +plan = source.prepare( + local=[source_registration.region()], + remote=[destination_peer.region("kv")], +) +work = source.write(plan, notification=b"kv-visible") +work.wait() + +work.close() +plan.close() +destination_peer.close() +source_registration.close() +destination_registration.close() +source.close() +destination.close() +``` + +External NIXL registration is explicit and lazy: + +```python +from nixl.torch_transfer import register_torch_backend + +register_torch_backend() +endpoint = _transfer.Endpoint("worker", backend="nixl") +``` + +Importing `torch` or `nixl` alone does not activate the integration. + +## Provider construction ownership + +Existing providers remain source-compatible with the version-1 factory, which +returns a `Backend`. New native providers should register FactoryV2 explicitly: + +```python +def factory(endpoint_id, name, incarnation, progress_mode, thread_mode, options, + *, adopt_backend): + backend = ProviderBackend.new_unactivated(...) + adopt_backend(backend) # Core owns cleanup from this point + backend.activate_native() # may publish native identity or start threads + return None + +register_backend("provider", factory, factory_api_version=2) +``` + +A version-1 factory return is adopted only after Core verifies that it is a +`Backend`. A wrong-type return is foreign and unowned: Core raises `TypeError` +without invoking `close()` or any other method on that object. Version 1 cannot +close the interruption window between the factory return and this adoption; +new providers should use FactoryV2. + +Core advertises `BACKEND_FACTORY_API_VERSION` as the maximum backward-compatible +factory contract it understands. It never inspects a callable signature or +retries under another version. A version-2 provider must adopt exactly once +before activation and return `None`. Core retains an interruption-durable close +receipt and weak finalizer; a failed orphan is quarantined by exact +`(backend_name, endpoint_id)`, so an unresolved identity collision fails closed +without blocking unrelated endpoint IDs. + +Backend registration itself is retry-safe across a lost return. Re-registering +the same name with the identical callable object and exact factory API version +is an idempotent no-op under the registry lock. A different callable identity +or version still raises `BusyError`; providers cannot use retry recovery to +replace another owner's registration. + +`Endpoint.close()` does not call `progress()` implicitly. Applications using +manual progress own their drain policy; Core still rejects active Work and open +child resources before it starts retry-only backend teardown. + +## Memory model + +`Endpoint.register` accepts a contiguous dense Tensor, Storage, writable +C-contiguous buffer, or capability-gated `RawSpan`. It registers the exact +logical byte interval and never creates a contiguous staging copy. A Tensor +registration anchors its original Storage, so rebinding the Tensor with `set_` +does not invalidate the registered address. + +The caller must keep the underlying allocation address and extent unchanged +until the Registration is closed. In particular, Storage must not be resized +while a plan, Work, or TransferSlot can reference it. Core checks the anchored +address and size +before metadata export, plan creation, and submission, but that check is only a +diagnostic: Python cannot close the race with a concurrent resize after work is +submitted. Overlapping local registrations are rejected, including aliases of +the same Tensor, Storage, or writable buffer. + +Named `Peer.region` is the safe remote API. `Peer.unsafe_region` exists only for +legacy protocols that exchange addresses; the provider must prove that the raw +span is covered by imported remote metadata. Backends that cannot prove this +must reject it. + +A region can represent one interval or a compact strided layout: + +```python +region = registration.region( + offset=0, + nbytes=block_bytes, + stride=block_stride, + count=num_blocks, +) +``` + +The allocation extent is bounds checked; `nbytes * count` is the transfer size. + +## Indexed plans + +`prepare(..., indexed=True)` creates reusable local and remote descriptor +catalogs. Each Region contributes `count` logical blocks, or one block when +`count` is absent. Indices address the flattened region-major, block-minor +catalog, so a compact strided Region can expose many independently selectable +blocks without expanding the plan input. `TransferPlan.local_block_count` and +`remote_block_count` report the resulting bounds. + +Local and remote block counts and aggregate sizes may differ. Each submission +selects independent local and remote block-index sequences, whose selected byte +totals must match. Each selected block contributes its Region's `nbytes` value. +Only selected destination blocks participate in alias validation; duplicate or +physically overlapping destination selections are rejected. Omitting both index +sequences selects every logical block in each catalog. + +```python +plan = endpoint.prepare( + local=[local_registration.region(0, block_bytes, stride=64, count=8)], + remote=[remote_peer.region("blocks", 0, block_bytes, stride=64, count=32)], + indexed=True, +) +work = endpoint.write( + plan, + local_indices=[7, 2], + remote_indices=[19, 4], +) +``` + +Individual providers may impose a stricter mapping. The current NIXL provider +requires selected local and remote sequences to have equal lengths and matching +per-block sizes. + +### One-shot and reusable dynamic index selections + +Frameworks that produce a new packed block table per request should use the +one-shot plan API. Core snapshots and authenticates the indices, but does not +allocate a transient public `IndexSelection` or enter the endpoint state lock a +second time: + +```python +import numpy as np + +work, state = plan.submit_indices_and_poll( + _transfer.TransferOp.WRITE, + local_indices=np.asarray(local_block_ids, dtype=np.int32), + remote_indices=np.asarray(remote_block_ids, dtype=np.int32), + notification=request_tag, + max_polls=0, +) +``` + +`plan.submit_indices(...)` is the post-only form. Providers without fused +submission support use the same provider-neutral fallback; framework adapters +do not need a capability branch. + +When the same validated selection will be bound, reused, or turned into a +persistent slot, retain it explicitly: + +```python +selection = plan.select_indices( + _transfer.TransferOp.WRITE, + local_indices=np.asarray(local_block_ids, dtype=np.int32), + remote_indices=np.asarray(remote_block_ids, dtype=np.int32), +) +work, state = selection.submit_and_poll( + notification=request_tag, + max_polls=64, + timeout_ns=None, +) +while state not in { + _transfer.WorkState.COMPLETED, + _transfer.WorkState.FAILED, + _transfer.WorkState.CANCELLED, +}: + state = work.poll_bounded( + max_polls=4096, + timeout_ns=100_000, + timeout_check_interval=32, + ) +work.wait() # raises for FAILED/CANCELLED, otherwise returns immediately +work.close() +``` + +Both `submit()` forms remain strictly post-only and are the correct QD>1 +primitive. A +provider may advertise `fused_prevalidated_submit_poll` and `fused_work_poll` +independently. The NIXL provider uses them to keep each bounded loop behind one +GIL release; providers without either capability use the source-compatible +composition. A zero initial poll budget posts only. On a non-fused provider its +returned Boolean is conservatively false because Core cannot recover a +post-derived native state without issuing a status observation. Boolean +`submit_indices_and_test`, `submit_and_test`, and `test_bounded` wrappers +remain available when a caller needs only terminal versus nonterminal. + +All bounded-poll counts, nanosecond budgets, and timeout-check intervals must +fit in a non-negative (or positive, where zero is not meaningful) signed +64-bit integer, matching the native provider ABI. A zero timeout still permits +exactly the first state observation even when the sampling interval is larger +than one; later clock checks occur at interval multiples. + +`IndexSelection` is immutable and belongs to the exact plan and operation used +to create it. Core copies both mutable, one-dimensional, C-contiguous native +`int32` inputs into one private packed buffer; NIXL reads that snapshot through +`numpy.frombuffer` without another copy. A general `Sequence` is still accepted: +Core applies `operator.index`, rejects booleans, validates bounds, and packs it +in one pass. Because this provider-neutral fast representation is signed +32-bit, an indexed catalog used through `select_indices` is limited to +2,147,483,648 logical entries. The format is deliberately native-endian so it +can be consumed directly by Torch and native providers; non-native-endian or +non-contiguous buffers take the safe sequence path. + +Uniform block size and full-catalog destination disjointness are plan facts. +Core can therefore validate byte totals without expanding regions and validate +destination uniqueness with one packed sort. Irregular or conservatively +ambiguous catalogs retain exact region expansion and physical-overlap checks. +The provider may omit its duplicate per-selection work only when Core's +authenticated token proves uniform geometry, full-catalog disjointness, and +unique destination indices *and* the provider independently proves the matching +full-catalog facts. + +Bind a selection when the same mapping is submitted repeatedly with a +request-local notification: + +```python +bound = selection.bind() +for request_tag in generation_tags: + work = bound.submit(notification=request_tag) + work.wait() + work.close() +bound.close() +``` + +The current NIXL raw binding updates or explicitly clears the notification at +every post. The Core reuse token identifies only the exact plan, operation, and +index snapshot; it never encodes provider-specific notification behavior. +Older NIXL bindings remain safe because the provider retires a cached request +before crossing between notification-present and notification-absent states. +Changing one non-empty tag to another still reuses the same request. Closing the +plan invalidates an otherwise live selection; each open Work or +selection-backed binding retains its packed snapshot until native access is +finished. + +## Bound submissions + +Latency-sensitive paths can bind one fully specified submission shape once and +reuse it without repeating plan ownership, storage stability, index, overlap, +notification, or CUDA-token validation: + +```python +bound = plan.bind( + _transfer.TransferOp.WRITE, + local_indices=[7, 2], + remote_indices=[19, 4], + cuda_ordering=_transfer.CudaOrderingMode.CALLER_READY, +) + +for _ in range(iterations): + work = bound.submit() + work.wait() + work.close() + +bound.close() +``` + +Every `submit()` creates a normal, independently owned `Work`; concurrent open +Work objects are permitted. Core still performs constant-time endpoint, +binding, and plan lifecycle checks and publishes all lifetime reservations +before calling a provider. For `host_wait`, static device and token validation +happens at bind time, while capture detection, event recording, and the actual +host wait happen on every submission. + +When a provider reports `FAILED`, Core derives and publishes the failure +diagnostic before committing the terminal Work state. If an asynchronous +`BaseException` interrupts that diagnostic lookup, Work remains nonterminal and +the next observation retries it; a terminal state can therefore never make a +missing diagnostic permanently unobservable. + +An ordinary `plan.bind(...)` binding remains fully immutable: operation, +indices, notification, CUDA ordering mode, stream, and ready event cannot be +overridden by `submit()`. The narrower `IndexSelection.bind()` form above owns +an immutable operation and index snapshot but intentionally accepts a +per-submission notification. Its provider path distinguishes a missing payload +from a replacement payload and cannot retain a stale tag. + +Providers receive an opaque `reuse_token`. Its identity is stable across all +submissions from one binding, allowing a provider to cache immutable native +lowering. An `IndexSelection` allocates its shared geometry token lazily on the +first selection-backed `BoundTransfer.submit()`: direct one-shot submissions +carry no token, and a binding used only to create a persistent request slot +allocates none. Selection-backed bindings also avoid a redundant private +binding token. The full-selection, notification-free, `caller_ready` +convenience forms `endpoint.read(plan)` and `endpoint.write(plan)` use stable +tokens per plan and operation. Tokens are hints only and never define transfer +semantics. + +### Persistent request slots + +A provider can advertise `Capabilities.persistent_request_slots` to expose the +lowest-allocation single-flight form of a bound transfer. A slot is created once +and owns one stable provider request; starting a generation does not allocate a +public `Work` or modify Core's resource sets: + +```python +bound = plan.bind( + _transfer.TransferOp.WRITE, + local_indices=[7, 2], + remote_indices=[19, 4], +) +slot = bound.create_slot() + +for _ in range(iterations): + slot.start_and_wait() + slot.recycle() + +slot.close() +bound.close() +``` + +For a blocking full-cycle loop, `execute()` combines start, completion, and +successful logical recycle under the slot's serialization boundary: + +```python +for _ in range(iterations): + completed_generation = slot.execute() +``` + +For queue-depth loops, `start_and_test()` exposes the initial post result with +zero additional status observations when a fused provider supports that exact +contract, while preserving Core ownership. It returns `True` with the slot logically +idle on immediate success, or `False` with an active generation. Poll pending +members with `test_and_recycle()`; it likewise returns `True` only after clean +successful completion has made the slot repostable. Failure and cancellation +are raised and remain owned for explicit recovery. Request-local notification +slots require `notification=` on `start_and_test`, exactly as on `start` and +`execute`. + +It returns the completed generation number. On success, `active` is false, +`state` is `None`, and the next generation or `close()` may proceed immediately. +Failure, cancellation, and timeout still leave the generation active, exactly +like `start_and_wait()`, so the caller can inspect or cancel it and must recycle +it explicitly after terminal state. The successful idle publication is atomic +with terminal observation under `SERIALIZED` and `MULTIPLE`; `SINGLE` provides +the same guarantee through its owner-thread contract, and +`CALLER_SERIALIZED` through the caller's whole-resource-graph exclusion. + +Providers may advertise +`Capabilities.deferred_successful_request_slot_recycle`. Core then omits the +provider `recycle()` call after clean `COMPLETED` in `execute()`, +`start_and_test()`, or `test_and_recycle()`. The provider retains that +observable completion receipt and consumes it inside the next `start()`, +`start_and_poll()`, or `close()` boundary. Explicit +`start()`/`wait()`/`start_and_wait()` behavior is unchanged and still requires +explicit `recycle()`. Providers without the capability remain source-compatible: +Core performs provider recycle before publishing execute success. + +Explicit `recycle()` completes provider recycle before making `active` false, +the sole public idle commit. Private cleanup follows, with `state = None` last +as a clean-idle marker. Public state and error ignore private terminal +diagnostics while idle; if asynchronous interruption stops cleanup, the next +begin observes the marker and idempotently repairs it before publishing another +generation. + +The exact `timeout=None` + lock-free + fused-start + deferred-success case is a +straight-line Core specialization. It applies to `SINGLE` and correctly used +`CALLER_SERIALIZED` endpoints and performs the same closed/creation, +conditional owner-thread, active-generation, CUDA-ordering, +generation-publication, and ambiguous-launch checks without crossing the +generic helper frames. Immediate clean completion publishes idle directly. +Because `_active` is its sole Core publication commit and the provider owns +durable state after dispatch entry, this exact synchronous path needs no +transient Core PENDING or success-side state-clear store; pre-entry interruption +becomes a retained Core-only failure. Nonterminal and exceptional results rejoin +the ordinary bounded/observed recovery paths. This is an internal latency +optimization only and does not change the public or provider contract. + +`TransferSlot` is intentionally mutable and single-flight. Its `state` is +`None` while idle. `start()` creates one owned generation, increments +`generation`, and changes `active` to true. On the explicit lifecycle, a second +start is rejected until the first generation reaches `COMPLETED`, `FAILED`, or +`CANCELLED` and the caller invokes `recycle()`. `active` remains true in terminal +state until that recycle boundary. Every `wait()` and `start_and_wait()` pins +the generation it began observing; if another thread recycles and restarts the +mutable slot before the waiter makes its next observation, that waiter raises +`BusyError` instead of reporting the replacement generation. A slot can close +only while idle. + +The ordinary binding and `TransferSlot` remain fully immutable: operation, +indices, attached notification, stream/event, and CUDA ordering cannot change +between generations. With a bound stream in `host_wait`, each slot owns one +reusable, lazily materialized CUDA event. Core caches its record/synchronize +callables and the stream capture query; every `start()` still checks capture, +records that same event after prior producer work, and host-waits before entering +the provider. An explicitly supplied ready event keeps its cached synchronize +callable. No readiness event is created in the steady generation path. + +When a framework reuses one exact operation/descriptor/index selection but +changes its completion tag, it may instead call +`binding.create_slot(request_local_notifications=True)`. The returned +`RequestLocalNotificationSlot` requires `notification=` on every `start`, +`start_and_wait`, or `execute`; `None` explicitly clears the native tag. This +mode is separately capability-gated so the fixed `TransferSlot.execute()` path +does not gain a request-local branch. The operation, selected indices, CUDA +ordering, and all memory ownership remain immutable. + +`start_and_wait(timeout)` is the latency-oriented equivalent of `start()` plus +`wait(timeout)`. A provider may additionally advertise +`Capabilities.fused_request_slot_start_poll` and implement the bounded +`BackendRequestSlot.start_and_poll(max_polls=..., timeout_ns=...)` SPI. Core +then keeps native post and up to 64 initial status observations inside one +provider boundary. The same SPI accepts `max_polls=0` for post-only execution: +an overriding fused provider returns the post-derived state, while the default +source composition can only return its cached `state`. A provider may independently advertise +`Capabilities.fused_request_slot_poll` and implement +`BackendRequestSlot.poll_bounded(max_polls=..., timeout_ns=..., +timeout_check_interval=...)`. If the request +is still pending or running, Core continues in chunks of at most 4,096 native +observations and a 100-microsecond soft residency target. The default interval +checks the timeout after observation 1 and then every 32 observations; one +provider call is non-preemptible and the sampled target can overshoot. Every +chunk reacquires the slot lock, verifies the pinned generation, invokes the +provider, publishes monotonic state, and releases the lock. This removes +per-observation Python crossings and most clock-read tax while retaining +count-bounded cancellation and signal checkpoints between chunks. Providers +without poll-only fusion retain the latency-oriented single-poll spin path; +providers without start fusion may still accelerate polling after ordinary +`start()`. + +At slot adoption, a provider may pass +`adopt_slot(wrapper, execution_dispatch=dispatch)`. Core strongly retains the +dispatch object and validates and caches its `start`, `poll_state`, `cancel`, +and `recycle` callables, plus `start_and_poll` and `poll_bounded` when their +capabilities are advertised. The wrapper remains authoritative for `error`, +`completion_event`, and `close`. Omitting `execution_dispatch` caches the +wrapper methods exactly as before, so existing providers remain source +compatible. The hook lets a native object own latency-sensitive execution +state without weakening Core's generation, timeout, CUDA-ordering, or lifetime +authority. + +For the deferred-success, fused, `SINGLE` `execute()` specialization, adoption +caches the eligible fused callable itself and whether CUDA readiness is already +caller-owned. The steady path branches on that callable before loading the +generic lock or dispatch table; serialized, timed, non-fused, and non-deferred +slots continue through the general path. + +Request-slot operations are interruption-safe but not signal-handler-reentrant: +after a Python signal interrupts an operation, let that frame unwind before +calling the same slot again. A subsequent ordinary call recovers the durable +state. This avoids charging every hot dispatch for a recursive-entry guard that +cannot make arbitrary signal-handler use safe. + +The standalone `wait()` operation retains cooperative yield/backoff for callers +that prefer lower host-CPU pressure; both operations have identical state, +timeout, cancellation, and lifetime semantics. + +Core passes the provider timeout as integer nanoseconds and saturates extremely +large finite public timeouts at signed 64-bit maximum for that bounded call; +the subsequent Core deadline remains an arbitrary-precision Python integer. +The private, introspectable policy constants are +`_FUSED_REQUEST_SLOT_POLL_MAX_POLLS` and +`_FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS`, and +`_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL`; benchmarks should read them rather than +duplicate the current 4,096/100-microsecond policy. + +The timeout is a total host-side budget from entry to `start_and_wait`. It +cannot preempt one blocking CUDA-readiness or provider call. A timeout never +recycles implicitly: the generation remains active and cancellable, and the +caller must observe a terminal state before explicitly invoking `recycle()`. +Every generation, fused or not, repeats `host_wait` readiness before Core +publishes the active generation and enters the provider. + +`completion_event` is exposed only when the backend advertises +`cuda_completion_event`, and only while a generation is active. The event is +generation-local: enqueue any device wait before `recycle()`, because a later +generation may re-record or replace the provider event. + +The provider slot is adopted only after Core has published the slot in the +complete endpoint/binding/plan lifetime graph and before the provider may create +native state. Once `start()` crosses into the provider, Core conservatively owns +that generation even if an asynchronous exception crosses the call boundary. +Providers must convert failures after possible native launch into terminal slot +state. The retained generation can then be inspected, cancelled where possible, +and recycled without losing native ownership. + +Slot creation publishes the Endpoint recovery edge before the binding blocker, +so an interruption cannot leave an inaccessible child that permanently blocks +its parent. Close uses closure as a retryable commit marker: it retains at least +one parent edge until that marker is visible, and a repeated `close()` finishes +any interrupted Core-only unlink/cleanup epilogue. Generation publication and +terminal failure/cancellation publication use the same commit-last rule. A +pre-provider interruption can therefore be recycled as a Core-only failed +generation without issuing provider `recycle()` against an idle native slot; +an interrupted terminal-error lookup is retried on the next observation. + +This opt-in path is for repeated immutable transfers. `BoundTransfer.submit()` +remains the general concurrent path and continues to return an independent +`Work` for every submission. A TransferSlot never makes `Work` reusable. + +## Work and ownership + +`Work.test()` is request-local and returns whether the request is terminal. +`Work.state` distinguishes completed, failed, and cancelled work; `Work.error` +retains the failure. `wait(timeout)` raises typed timeout, cancellation, or +backend errors. Cancellation is best effort and does not roll back remote writes. + +Core retains the plan, registrations, peer, ordering tokens, and notification +payload until Work is closed. Closing a live or referenced resource raises +`BusyError`. Close operations are idempotent after success. Required teardown is: + +```text +Work -> BoundTransfer (when used) -> TransferPlan -> Peer / Registration -> Endpoint +TransferSlot -> BoundTransfer -> TransferPlan -> Peer / Registration -> Endpoint +``` + +`Endpoint.open_bound_transfers` is the recovery snapshot for bindings. A +binding cannot close while it has an open Work or TransferSlot, a plan cannot +close while it has an open binding, and an endpoint cannot close until all +bindings are closed. `Endpoint.open_transfer_slots` exposes idle, in-flight, +terminal, and cleanup-only slots for recovery after interruption. + +One-shot READ/WRITE owns an internal plan and releases it when Work closes. +The provider SPI uses an explicit adoption handshake for data transfers and +standalone notifications: Core publishes the Work and its complete lifetime +graph first, then the provider must adopt a `BackendWork` before launching +native activity. A provider may raise before adoption only when no native work +remains; after adoption, launch errors must become terminal Work failures. This +closes the interruption window between native submission and Python +bookkeeping. `Endpoint.open_works` provides a recovery snapshot if an +asynchronous exception lands after submission but before caller code stores the +returned Work. + +Code that must close that final caller-side adoption window may pass a unique +opaque `recovery_token=` to any Work-producing operation. Core stores the +object by identity before publishing the Work and exposes it unchanged through +`Work.recovery_token`. Recovery is deliberately a cold-path scan: + +```python +attempt = object() +try: + work = plan.submit_indices( + "write", + local_indices=local_indices, + remote_indices=remote_indices, + recovery_token=attempt, + ) +except BaseException: + matches = [ + candidate + for candidate in endpoint.open_works + if candidate.recovery_token is attempt + ] + # zero: Core retained no provider activity; one: the lost Work; + # more than one: an ambiguous token-uniqueness error. + raise +``` + +Core never hashes the token, compares it, coerces it to Boolean, invokes it, +serializes it, or passes it to the provider. It is held strongly until +`Work.close()` removes every lifetime edge, then cleared as the close +epilogue's last commit. Use `is`, never `==`, and do not use the default `None` +for recovery matching. This is distinct from provider-facing `reuse_token`: +the reuse token identifies a cached native request shape; the recovery token +identifies one caller submission attempt. + +Registration, peer import, and plan preparation use the same ownership rule: +Core publishes a resource wrapper, the provider adopts a retry-safe transaction +record before its first native mutation, and Core rolls it back on failure. If +rollback itself fails, the wrapper is cleanup-only and remains recoverable via +`Endpoint.open_registrations`, `open_peers`, or `open_plans`; `close()` retries +the provider cleanup. + +Endpoint construction is the one boundary where Core cannot first publish a +recovery object: Python has not returned the `Endpoint` to its caller yet. +Provider factories must therefore catch `BaseException` and self-roll back all +native state created before they return. This is a provider obligation, not an +exactly-once guarantee that Core can enforce around arbitrary third-party +factory code. + +## Progress and threading + +`progress_mode` is `manual` or `background`. Background progress is provider +owned; Core does not create a hidden thread. `thread_mode` is `single`, +`caller_serialized`, `serialized`, or `multiple`: + +- `single` is confined to the thread that created the Endpoint and has no Core + locks; +- `caller_serialized` may migrate between threads, but the application must + prevent overlap and reentry for the entire Endpoint resource graph, including + Work/slot observation, progress, cancellation, snapshots, waits, and teardown; +- `serialized` lets calls migrate and makes Core provide that exclusion, + including provider entry; +- `multiple` permits concurrent operations when the provider advertises that + capability. Core publishes Work/lifetime reservations before releasing its + lifecycle lock, so data submission and status, standalone notification send, + destructive notification receive, and manual progress can overlap when their + backend resources are independent. Destructive notification consumers remain + serialized with each other. Endpoint close returns `BusyError` while an + admitted provider call is active and never tears down its target concurrently. + Providers still serialize private destructive state that cannot safely join + that overlap. + +Provider-owned background progress threads are outside the caller boundary. +An application must keep its external exclusion for the complete public call, +including a blocking `host_wait` or `await Work` iteration. Selecting +`caller_serialized` without that whole-graph contract is unsafe. + +Manual progress is endpoint-local. Every participating process must call +`Endpoint.progress()` regularly, including an otherwise idle remote responder; +waiting on one endpoint cannot drive a different process's transport engine. + +## Device ordering + +READ and WRITE expose two explicit CUDA producer-ordering modes: + +- `caller_ready` (default): the caller guarantees that selected local memory is + ready before submission. No CUDA API is called. +- `host_wait`: pass exactly one real `torch.cuda.Stream` or recorded + `torch.cuda.Event`. Core validates that every selected local Region is on one + explicitly indexed CUDA device, records one event when a stream is supplied, + then waits for that event on the host before posting the provider request. + +The wait happens outside Core's state and provider locks. It never inserts a +device-wide synchronization, and providers must not add another implicit wait. +`host_wait` protects local memory only: readiness of a remote READ source still +belongs to the remote application's control protocol. Core checks capture on a +supplied stream before recording its event. A supplied ready event goes directly +to its authoritative `synchronize()` operation and does not query the unrelated +current stream's capture state; an event CUDA cannot synchronize still fails +through CUDA. +Native stream-attached, graph-safe transport submission would require a future, +separately advertised mode. + +Terminal Work completion means the local allocation can be reused. A provider +may additionally advertise `cuda_completion_event`, in which case +`Work.completion_event` expresses that terminal point to CUDA. The initial NIXL +provider does not advertise such an event. + +## Notifications + +Capabilities distinguish two contracts that native integrations use +differently: + +- `attached_notifications`: a notification attached to READ/WRITE is + transfer-ordered and cannot become observable before the associated remote + writes are visible. +- `standalone_notifications`: `send_notification` has the terminal meaning in + `standalone_notification_completion`. NIXL advertises `local_accepted`: a + completed Work means the native API accepted/copied the payload, not that the + remote application observed it. Remote delivery still requires transport + progress and receiver polling. + +A backend may additionally advertise `synchronous_notification_send` and +implement `Endpoint.send_notification_sync(peer, payload)`. This additive fast +path returns `None` and creates no Core or provider `Work`; return still means +only `standalone_notification_completion`. For NIXL that is `local_accepted`, +not remote observation. An exception after provider entry can therefore be +outcome-ambiguous. Reliable control protocols must retain sequence IDs, +acknowledgement, deduplication, and retry. Callers that need inspectable local +ownership after interruption should keep using `send_notification()`. +In `multiple` mode Core uses a short peer-lifetime reservation and releases its +endpoint lifecycle lock before provider entry, so independent synchronous sends +retain provider/native concurrency. The reservation is a pooled strong token in +the send frame plus a close-visible weak identity: steady state allocates no new +token, and an interrupted cleanup leaves a dead identity that `Peer.close()` +prunes rather than a permanent false blocker. `Peer.close()` reports +`BusyError` only while a live reservation exists. Providers rely on this Core +pin and should not duplicate it on their hot SPI path. `serialized` continues +to exclude the complete call, while `single` and `caller_serialized` retain +their caller contracts. + +The asynchronous `send_notification()` path applies the same lock-breadth rule +without another token allocation: its required `Work` is published with the +peer and payload lifetime anchors before provider entry. Progress and receive +use pooled close-visible call reservations, allocated only when concurrency +exceeds the warmed high-water mark. Receive calls take a dedicated consumer +lock because dequeue is destructive; that lock does not exclude data submit, +independent Work status, asynchronous sends, or provider-safe progress. + +Received notifications carry either no source identity or the complete +`(endpoint_id, incarnation)` pair. Consumers must reject stale or unknown +incarnations rather than routing by a reusable display name alone. + +`poll_notifications()` is an at-most-once receive boundary. The current SPI +allows a provider to dequeue before returning the batch, so an asynchronous +exception or allocation failure after dequeue can lose that batch. Protocols +that require reliable control messages must include an application sequence ID, +acknowledgement, deduplication, and retransmission. A future peek/token/ack SPI +would be required to provide a stronger receive transaction. + +High-rate consumers should use +`poll_notification_batches(max_items=None)`. `max_items` counts payloads across +all returned batches; zero returns without entering the provider. Each +`NotificationBatch` contains a nonempty immutable `tuple[bytes, ...]` and either +a complete source identity or no identity. Repeated batches for one source are +legal and preserve that source's observation order. Batch boundaries have no +wire meaning, and no total order is promised between different sources. + +Providers may advertise `grouped_notification_receive` when they retain their +native source grouping instead of first constructing one provider object per +payload. The public operation remains available without that capability through +a scalar compatibility implementation. It groups only adjacent notifications +with the same complete identity and keeps anonymous notifications separate. +Payload tuples that are already exact immutable bytes are retained by identity; +mutable inputs are snapshotted. Returned-value validation and public object +construction run after Core releases its endpoint state lock, while the +destructive provider call itself remains serialized with close and mutation. + +`max_items` is an egress bound, not automatically a native-ingress memory bound. +A provider whose native API drains a complete transport batch before returning +must separately cap or document that temporary materialization. The current +NIXL provider bounds retained Python count and bytes and fails closed on +malformed/overflowed ingress, but the underlying NIXL `getNotifs` API does not +yet offer a bounded dequeue. + +## Metadata and security boundary + +`EndpointMetadata` uses bounded canonical JSON with base64 opaque provider +fields, a `{major, minor}` wire version, required-feature names, optional opaque +extensions, stable endpoint ID, and restart incarnation. Unknown outer fields, +missing or duplicate fields, unknown required features, malformed values, and +unsupported major versions are rejected. A higher minor version is accepted +only through understood required features and the extension container. Encoded +metadata, provider payloads, descriptors, registration counts, notification +payloads, and identity strings have explicit limits. Opaque capability-bearing +provider data is redacted from object representations. No pickle is used by +Core. + +Callers with a stable logical identity may pass `endpoint_id=`. Core creates a +fresh incarnation for each Endpoint construction and rejects replacement +metadata until the old Peer is explicitly closed. + +When the control plane already knows a peer identity, use +`import_peer(expected_identity=...)`, or the individual `expected_endpoint_id`, +`expected_name`, and `expected_incarnation` checks. Core performs these checks +before the provider can mutate native peer state. + +`export_metadata(registrations=[...])` is a capability grant for exactly the +selected registrations. Providers must emit partial transport metadata or reject +the operation; silently exposing omitted registrations is not allowed. +In this initial contract each imported Peer is an immutable grant snapshot. +Extending or revoking that grant requires closing dependent plans and the Peer, +then importing a replacement snapshot. Incremental `Peer.update` semantics are +left open because merge, revocation, and in-flight-work ordering must be defined +together. + +Metadata authenticity, confidentiality, replay protection, and application-level +authorization remain the control plane's responsibility. Metadata is a bearer +capability: do not send it over an unauthenticated channel or log its opaque +fields. Core's bounds and strict parser are defense in depth, not a secure +transport envelope. + +The exporting endpoint cannot observe remote plans or work. The application +must therefore hold a registration lease until every peer has quiesced and +released work that may reference it, then communicate revocation before local +deregistration. Core's local lifetime anchors do not implement a distributed +lease protocol. + +## Initial limitations + +- private Python API and SPI; no stable ABI; +- CPU reference backend only for conformance; +- the NIXL provider has two-node GB200/UCX correctness and large-transfer + performance coverage, but not full framework end-to-end GPU coverage; +- no atomic model activation, remote lease protocol, or rollback; +- no FILE/OBJ tier, distributed-checkpoint durability, or NIXL-EP token path; +- `host_wait` is a host-side gate, not a native GPUNETIO stream submission; +- no CUDA Graph or compiler integration and no CUDA completion event from NIXL; +- no promise that cancellation prevents partial remote visibility. +- full-catalog overlap checks conservatively reject compact strided regions + whose bounding envelopes interleave; explicit indexed selections are checked + exactly without expanding the whole catalog. diff --git a/torch/distributed/_transfer/__init__.py b/torch/distributed/_transfer/__init__.py new file mode 100644 index 0000000000000..9c181f25dea2f --- /dev/null +++ b/torch/distributed/_transfer/__init__.py @@ -0,0 +1,123 @@ +"""Experimental backend-neutral endpoint transfer API. + +This module is private while its ownership, ordering, and provider contracts are +validated with real inference and training consumers. +""" + +from ._api import ( + BoundTransfer, + Endpoint, + IndexSelection, + Peer, + RawSpan, + RequestLocalNotificationSlot, + Region, + Registration, + TransferPlan, + TransferSlot, + Work, +) +from ._backend import ( + BACKEND_FACTORY_API_VERSION, + Backend, + BackendFactory, + BackendFactoryV2, + BackendIndexSelection, + BackendRequestSlot, + BackendWork, + ImportedRawMemory, + ImportedRegistration, + RegisteredMemory, + register_backend, + unregister_backend, +) +from ._errors import ( + BackendError, + BackendFailureError, + BackendMismatchError, + BackendNotFoundError, + BusyError, + ClosedError, + InvalidMetadataError, + InvalidRegionError, + TransferCancelledError, + TransferError, + TransferTimeoutError, + UnsupportedError, +) +from ._reference import _register_reference_backend +from ._types import ( + BackendNotification, + BackendNotificationBatch, + BackendRegion, + Capabilities, + CudaOrderingMode, + EndpointIdentity, + EndpointInfo, + EndpointMetadata, + Notification, + NotificationBatch, + NotificationCompletion, + ProgressMode, + RegistrationMetadata, + ThreadMode, + TransferOp, + WorkState, +) + + +_register_reference_backend() + +__all__ = [ + "BACKEND_FACTORY_API_VERSION", + "Backend", + "BackendError", + "BackendFactory", + "BackendFactoryV2", + "BackendIndexSelection", + "BackendFailureError", + "BackendMismatchError", + "BackendNotFoundError", + "BackendNotification", + "BackendNotificationBatch", + "BackendRegion", + "BackendRequestSlot", + "BackendWork", + "BoundTransfer", + "BusyError", + "Capabilities", + "ClosedError", + "CudaOrderingMode", + "Endpoint", + "EndpointIdentity", + "EndpointInfo", + "EndpointMetadata", + "ImportedRawMemory", + "ImportedRegistration", + "IndexSelection", + "InvalidMetadataError", + "InvalidRegionError", + "Notification", + "NotificationBatch", + "NotificationCompletion", + "Peer", + "ProgressMode", + "RawSpan", + "RequestLocalNotificationSlot", + "Region", + "RegisteredMemory", + "Registration", + "RegistrationMetadata", + "ThreadMode", + "TransferCancelledError", + "TransferError", + "TransferOp", + "TransferPlan", + "TransferSlot", + "TransferTimeoutError", + "UnsupportedError", + "Work", + "WorkState", + "register_backend", + "unregister_backend", +] diff --git a/torch/distributed/_transfer/_api.py b/torch/distributed/_transfer/_api.py new file mode 100644 index 0000000000000..06fde259e7a85 --- /dev/null +++ b/torch/distributed/_transfer/_api.py @@ -0,0 +1,8283 @@ +from __future__ import annotations + +import asyncio +import ctypes +import functools +import itertools +import math +import operator +import sys +import threading +import time +import uuid +import weakref +from bisect import bisect_right +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, cast, ParamSpec, TypeVar + +import torch +from typing_extensions import Self + +try: + from torch._C._distributed_transfer import ( # type: ignore[attr-defined] + _snapshot_and_analyze_int32_indices as _compiled_index_analyzer, + ) +except ImportError: + # Source-tree Python can run against an older installed torch._C while this + # private API is under development. Mixed/general sequences also never + # depend on the compiled helper. + _compiled_index_analyzer = None + +from ._backend import ( + Backend, + BackendIndexSelection, + BackendRequestSlot, + BackendWork, + ImportedRawMemory, + ImportedRegistration, + RegisteredMemory, + _get_backend_registration, + _make_backend_index_selection, +) +from ._errors import ( + BackendFailureError, + BackendMismatchError, + BusyError, + ClosedError, + InvalidMetadataError, + InvalidRegionError, + TransferCancelledError, + TransferError, + TransferTimeoutError, + UnsupportedError, +) +from ._types import ( + BackendNotification, + BackendNotificationBatch, + BackendRegion, + Capabilities, + CudaOrderingMode, + EndpointIdentity, + EndpointInfo, + EndpointMetadata, + Notification, + NotificationBatch, + ProgressMode, + RegistrationMetadata, + ThreadMode, + TransferOp, + WorkState, + _notification_batch_from_backend, +) + + +__all__ = [ + "BoundTransfer", + "Endpoint", + "IndexSelection", + "Peer", + "RawSpan", + "RequestLocalNotificationSlot", + "Region", + "Registration", + "TransferPlan", + "TransferSlot", + "Work", +] + +_TERMINAL_STATES = frozenset( + {WorkState.COMPLETED, WorkState.FAILED, WorkState.CANCELLED} +) +_CONSTRUCTION_TOKEN = object() +_UNADOPTED_HANDLE = object() +_UNSET_NOTIFICATION = object() +_MAX_NOTIFICATION_BYTES = 1024 * 1024 +_FUSED_REQUEST_SLOT_MAX_POLLS = 64 +_MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS = (1 << 63) - 1 +_FUSED_REQUEST_SLOT_POLL_MAX_POLLS = 4096 +_FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS = 100_000 +_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL = 32 +_GET_IDENT = threading.get_ident +_WORK_COMPLETED = WorkState.COMPLETED +_WORK_RUNNING = WorkState.RUNNING +_NATIVE_INT32_FORMATS = frozenset( + {"i", "@i", "=i", "i"} +) +_INT32_MAX = (1 << 31) - 1 +_INDEX_SELECTION_TORCH_THRESHOLD = 128 +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +class _BackendConstructionReceipt: + """Core ownership record independent of an Endpoint object's lifetime.""" + + __slots__ = ( + "_adoption_count", + "_backend", + "_closed", + "_lock", + "_poison_message", + "_published", + "_sealed", + "backend_name", + "endpoint_id", + ) + + def __init__(self, backend_name: str, endpoint_id: str) -> None: + self.backend_name = backend_name + self.endpoint_id = endpoint_id + self._lock = threading.RLock() + self._backend: Backend | None = None + self._adoption_count = 0 + self._sealed = False + self._published = False + self._closed = False + self._poison_message: str | None = None + + @property + def backend(self) -> Backend | None: + with self._lock: + return self._backend + + @property + def published(self) -> bool: + with self._lock: + return self._published + + @property + def closed(self) -> bool: + with self._lock: + return self._closed + + def adopt(self, backend: Backend) -> None: + if not isinstance(backend, Backend): + message = "backend factory must adopt a Backend instance" + with self._lock: + if self._poison_message is None: + self._poison_message = message + raise TypeError(message) + with self._lock: + if self._sealed: + message = "backend factory attempted adoption after its call returned" + if self._poison_message is None: + self._poison_message = message + raise BackendFailureError(message) + if self._backend is not None or self._adoption_count: + message = "backend factory must call adopt_backend exactly once" + if self._poison_message is None: + self._poison_message = message + raise BackendFailureError(message) + # Ownership commits before the diagnostic count. An interruption + # between these stores still leaves the native owner closeable. + self._backend = backend + self._adoption_count = 1 + + def seal_failed_call(self) -> None: + with self._lock: + self._sealed = True + + def seal_successful_call(self, result: object) -> Backend: + with self._lock: + self._sealed = True + if self._poison_message is not None: + raise BackendFailureError(self._poison_message) + if result is not None: + raise TypeError("version-2 backend factory must return None") + if self._adoption_count != 1 or self._backend is None: + raise BackendFailureError( + "version-2 backend factory must call adopt_backend exactly once" + ) + return self._backend + + def mark_published(self) -> None: + with self._lock: + if self._closed or self._backend is None: + raise BackendFailureError("cannot publish a closed backend receipt") + self._published = True + + def close(self, close_backend: Callable[[], object] | None = None) -> None: + """Close once; an interrupted/failed close remains retryable.""" + + with self._lock: + if self._closed: + return + backend = self._backend + if backend is None: + self._closed = True + return + callback = backend.close if close_backend is None else close_backend + try: + result = callback() + if result is not None: + raise BackendFailureError("backend close must return None") + except BaseException: + raise + self._closed = True + + +_ORPHANED_BACKEND_RECEIPTS: dict[tuple[str, str], set[_BackendConstructionReceipt]] = {} +_ORPHANED_BACKEND_RECEIPTS_LOCK = threading.Lock() + + +def _retain_orphaned_backend_receipt(receipt: _BackendConstructionReceipt) -> None: + key = (receipt.backend_name, receipt.endpoint_id) + with _ORPHANED_BACKEND_RECEIPTS_LOCK: + _ORPHANED_BACKEND_RECEIPTS.setdefault(key, set()).add(receipt) + + +def _forget_orphaned_backend_receipt(receipt: _BackendConstructionReceipt) -> None: + key = (receipt.backend_name, receipt.endpoint_id) + with _ORPHANED_BACKEND_RECEIPTS_LOCK: + receipts = _ORPHANED_BACKEND_RECEIPTS.get(key) + if receipts is None: + return + receipts.discard(receipt) + if not receipts: + _ORPHANED_BACKEND_RECEIPTS.pop(key, None) + + +def _finalize_backend_receipt(receipt: _BackendConstructionReceipt) -> None: + """Best-effort GC cleanup without retaining the Endpoint itself.""" + + if receipt.backend is None or receipt.closed: + return + _retain_orphaned_backend_receipt(receipt) + try: + # No Endpoint remains to supply its SERIALIZED lock. The receipt lock + # still serializes finalizer cleanup and all later orphan retries. + receipt.close() + except BaseException: + return + _forget_orphaned_backend_receipt(receipt) + + +def _retry_orphaned_backend_receipts(backend_name: str, endpoint_id: str) -> None: + """Audit one provider's orphans; only an exact identity blocks creation.""" + + with _ORPHANED_BACKEND_RECEIPTS_LOCK: + snapshot = tuple( + (key, tuple(receipts)) + for key, receipts in _ORPHANED_BACKEND_RECEIPTS.items() + if key[0] == backend_name + ) + + exact_failures: list[tuple[_BackendConstructionReceipt, Exception]] = [] + for key, receipts in snapshot: + for receipt in receipts: + try: + receipt.close() + except BaseException as error: + # Process-control interruptions are not converted or hidden; + # the retained receipt makes the next attempt safe. + if not isinstance(error, Exception): + raise + if key[1] == endpoint_id: + exact_failures.append((receipt, error)) + else: + _forget_orphaned_backend_receipt(receipt) + + if exact_failures: + receipt, error = exact_failures[0] + origin = ( + "an abandoned published endpoint (possibly with live provider children)" + if receipt.published + else "an unpublished endpoint construction" + ) + failure = BackendFailureError( + f"cannot reuse endpoint identity {endpoint_id!r} for backend " + f"{backend_name!r}: cleanup of {origin} is still pending" + ) + add_note = getattr(failure, "add_note", None) + if callable(add_note): + add_note(f"orphan cleanup failure type: {type(error).__name__}") + raise failure from error + + +class _NoopLock: + """Context-manager lock for SINGLE or externally serialized resources.""" + + __slots__ = () + + def __enter__(self) -> _NoopLock: + return self + + def __exit__( + self, + exc_type: object, + exc_value: object, + traceback: object, + ) -> bool: + return False + + +_NOOP_LOCK = _NoopLock() + + +class _ReuseToken: + """Weak-referenceable identity token with no provider-visible semantics.""" + + __slots__ = ("__weakref__",) + + +class _SyncNotificationReservation: + """Pooled strong token with a close-visible weak identity.""" + + __slots__ = ("active", "reference", "__weakref__") + + def __init__(self) -> None: + self.active = False + self.reference = weakref.ref(self) + + +class _BackendCallReservation: + """Reusable MULTIPLE-call token whose weak identity is visible to close.""" + + __slots__ = ("active", "reference", "__weakref__") + + def __init__(self) -> None: + self.active = False + self.reference = weakref.ref(self) + + +@dataclass(slots=True) +class _LifetimeAnchor: + owner: object + storage: object | None = None + buffer: memoryview | None = None + allocation_address: int | None = None + allocation_nbytes: int | None = None + + +@dataclass(frozen=True, slots=True) +class _RemoteRawDescriptor: + address: int + nbytes: int + device: str + memory_type: str + + +def _state_locked(function: Callable[_P, _R]) -> Callable[_P, _R]: + """Serialize endpoint resource-graph mutations.""" + + @functools.wraps(function) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + endpoint = cast(Any, args[0]) + state_lock = endpoint._state_lock + if state_lock is _NOOP_LOCK: + return function(*args, **kwargs) + with state_lock: + return function(*args, **kwargs) + + return wrapper + + +@dataclass(frozen=True, slots=True) +class RawSpan: + """Expert-only local address range with an explicit lifetime anchor.""" + + address: int + nbytes: int + device: str + memory_type: str + owner: object = field(repr=False) + + def __post_init__(self) -> None: + _validate_positive_int(self.address, "address") + _validate_positive_int(self.nbytes, "nbytes") + _validate_nonempty_str(self.device, "device", InvalidRegionError) + _validate_nonempty_str(self.memory_type, "memory_type", InvalidRegionError) + if self.owner is None: + raise InvalidRegionError("a local RawSpan requires a lifetime owner") + + +class Region: + """A bounded local or remote byte region. + + Regions are created by :meth:`Registration.region`, :meth:`Peer.region`, + or :meth:`Peer.unsafe_region`; users do not construct them directly. + """ + + __slots__ = ( + "_backend_descriptor", + "_device", + "_memory_type", + "_nbytes", + "_offset", + "_origin", + "_remote", + "_stride", + "_count", + ) + + def __init__( + self, + origin: Registration | Peer, + backend_descriptor: object, + *, + offset: int, + nbytes: int, + device: str, + memory_type: str, + remote: bool, + stride: int | None, + count: int | None, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("Region objects are created by Registration or Peer") + self._origin = origin + self._backend_descriptor = backend_descriptor + self._offset = offset + self._nbytes = nbytes + self._device = device + self._memory_type = memory_type + self._remote = remote + self._stride = stride + self._count = count + + @property + def offset(self) -> int: + return self._offset + + @property + def nbytes(self) -> int: + """Number of bytes in each strided segment.""" + + return self._nbytes + + @property + def device(self) -> str: + return self._device + + @property + def memory_type(self) -> str: + return self._memory_type + + @property + def stride(self) -> int | None: + return self._stride + + @property + def count(self) -> int | None: + return self._count + + @property + def total_nbytes(self) -> int: + return self._nbytes * (self._count or 1) + + @property + def is_remote(self) -> bool: + return self._remote + + def _to_backend(self) -> BackendRegion: + self._origin._check_open() + descriptor = self._backend_descriptor + if isinstance(self._origin, Registration): + descriptor = self._origin._handle + elif isinstance(descriptor, RegistrationMetadata): + descriptor = ImportedRegistration( + peer_handle=self._origin._handle, + registration_id=descriptor.registration_id, + name=descriptor.name, + nbytes=descriptor.nbytes, + device=descriptor.device, + memory_type=descriptor.memory_type, + descriptor=descriptor.backend_descriptor, + ) + elif isinstance(descriptor, _RemoteRawDescriptor): + descriptor = ImportedRawMemory( + peer_handle=self._origin._handle, + address=descriptor.address, + nbytes=descriptor.nbytes, + device=descriptor.device, + memory_type=descriptor.memory_type, + ) + return BackendRegion( + owner=self._origin, + offset=self._offset, + nbytes=self._nbytes, + device=self._device, + memory_type=self._memory_type, + backend_descriptor=descriptor, + remote=self._remote, + stride=self._stride, + count=self._count, + ) + + +class Registration: + """A backend registration that retains its source allocation. + + Retention does not freeze mutable Storage metadata. The caller must keep + the allocation address and extent unchanged until this object is closed; + in particular, resizing while a plan or Work exists is unsafe. Core's + pointer/size checks are diagnostics and cannot close the mutation race. + """ + + __slots__ = ( + "_closed", + "_creation_failed", + "_endpoint", + "_handle", + "_memory", + ) + + def __init__( + self, + endpoint: Endpoint, + memory: RegisteredMemory, + handle: object, + *, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("Registration objects are created by Endpoint") + self._endpoint = endpoint + self._memory = memory + self._handle = handle + self._creation_failed = False + self._closed = False + + @property + def registration_id(self) -> str: + return self._memory.registration_id + + @property + def name(self) -> str: + return self._memory.name + + @property + def nbytes(self) -> int: + return self._memory.nbytes + + @property + def device(self) -> str: + return self._memory.device + + @property + def memory_type(self) -> str: + return self._memory.memory_type + + @property + def closed(self) -> bool: + return self._closed + + def region( + self, + offset: int = 0, + nbytes: int | None = None, + *, + stride: int | None = None, + count: int | None = None, + ) -> Region: + self._endpoint._validate_registration(self) + offset, nbytes, stride, count = _validate_region_bounds( + self.nbytes, + offset, + nbytes, + stride, + count, + ) + if stride is not None and not self._endpoint.capabilities.strided_regions: + raise UnsupportedError( + f"backend {self._endpoint.backend!r} does not support strided regions" + ) + return Region( + self, + None, + offset=offset, + nbytes=nbytes, + device=self.device, + memory_type=self.memory_type, + remote=False, + stride=stride, + count=count, + _token=_CONSTRUCTION_TOKEN, + ) + + def close(self) -> None: + self._endpoint._close_registration(self) + + def _check_open(self) -> None: + if self._closed: + raise ClosedError(f"registration {self.name!r} is closed") + if self._creation_failed: + raise BackendFailureError( + f"registration {self.name!r} is retained for cleanup only" + ) + self._endpoint._check_open() + + def __enter__(self) -> Self: + self._check_open() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class Peer: + """An imported remote endpoint and its named registrations.""" + + __slots__ = ( + "_active_sync_notification_sends", + "_closed", + "_creation_failed", + "_endpoint", + "_endpoint_id", + "_handle", + "_incarnation", + "_name", + "_registrations", + "_sync_notification_reservation_pool", + ) + + def __init__( + self, + endpoint: Endpoint, + metadata: EndpointMetadata, + handle: object, + *, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("Peer objects are created by Endpoint.import_peer") + self._endpoint = endpoint + self._endpoint_id = metadata.endpoint_id + self._name = metadata.name + self._incarnation = metadata.incarnation + self._handle = handle + self._creation_failed = False + self._registrations = { + registration.name: registration for registration in metadata.registrations + } + # MULTIPLE-mode synchronous sends have no Work object to retain this + # peer. The send frame strongly owns a pooled token while this set keeps + # only its weak identity. If asynchronous interruption prevents the + # cleanup epilogue, the dead weak reference cannot wedge peer closure. + self._active_sync_notification_sends: set[ + weakref.ReferenceType[_SyncNotificationReservation] + ] = set() + self._sync_notification_reservation_pool: list[_SyncNotificationReservation] = ( + [] + ) + self._closed = False + + @property + def endpoint_id(self) -> str: + return self._endpoint_id + + @property + def name(self) -> str: + return self._name + + @property + def incarnation(self) -> str: + return self._incarnation + + @property + def identity(self) -> EndpointIdentity: + return EndpointIdentity(self.endpoint_id, self.incarnation) + + @property + def closed(self) -> bool: + return self._closed + + def region( + self, + name: str, + offset: int = 0, + nbytes: int | None = None, + *, + stride: int | None = None, + count: int | None = None, + ) -> Region: + self._check_open() + try: + metadata = self._registrations[name] + except KeyError as error: + raise InvalidRegionError( + f"peer {self.name!r} has no registration {name!r}" + ) from error + offset, nbytes, stride, count = _validate_region_bounds( + metadata.nbytes, + offset, + nbytes, + stride, + count, + ) + if stride is not None and not self._endpoint.capabilities.strided_regions: + raise UnsupportedError( + f"backend {self._endpoint.backend!r} does not support strided regions" + ) + return Region( + self, + metadata, + offset=offset, + nbytes=nbytes, + device=metadata.device, + memory_type=metadata.memory_type, + remote=True, + stride=stride, + count=count, + _token=_CONSTRUCTION_TOKEN, + ) + + def unsafe_region( + self, + address: int, + nbytes: int, + *, + device: str, + memory_type: str, + stride: int | None = None, + count: int | None = None, + ) -> Region: + """Create a provider-validated remote raw-address region. + + This escape hatch exists for legacy protocols that exchange addresses. + The backend must validate that the span is covered by imported metadata. + """ + + self._check_open() + if not self._endpoint.capabilities.raw_spans: + raise UnsupportedError( + f"backend {self._endpoint.backend!r} does not support raw spans" + ) + _validate_positive_int(address, "address") + _validate_positive_int(nbytes, "nbytes") + _validate_nonempty_str(device, "device", InvalidRegionError) + _validate_nonempty_str(memory_type, "memory_type", InvalidRegionError) + stride, count = _validate_stride(stride, count) + if stride is not None and stride < nbytes: + raise InvalidRegionError("stride must be at least nbytes") + descriptor = _RemoteRawDescriptor( + address=address, + nbytes=_region_extent(nbytes, stride, count), + device=device, + memory_type=memory_type, + ) + return Region( + self, + descriptor, + offset=0, + nbytes=nbytes, + device=device, + memory_type=memory_type, + remote=True, + stride=stride, + count=count, + _token=_CONSTRUCTION_TOKEN, + ) + + def close(self) -> None: + self._endpoint._close_peer(self) + + def _check_open(self) -> None: + if self._closed: + raise ClosedError(f"peer {self.name!r} is closed") + if self._creation_failed: + raise BackendFailureError( + f"peer {self.name!r} is retained for cleanup only" + ) + self._endpoint._check_open() + + def __enter__(self) -> Self: + self._check_open() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class TransferPlan: + """Reusable provider descriptors for one peer and one region layout. + + Indexed plans expose flattened region-major, logical-block catalogs. Each + Region contributes ``count`` blocks, or one block when ``count`` is absent. + """ + + __slots__ = ( + "_bindings", + "_closed", + "_creation_failed", + "_endpoint", + "_handle", + "_indexed", + "_local", + "_local_block_count", + "_local_full_nonoverlap", + "_local_starts", + "_local_total_nbytes", + "_local_uniform_device", + "_local_uniform_nbytes", + "_origins", + "_remote", + "_remote_block_count", + "_remote_full_nonoverlap", + "_remote_starts", + "_remote_total_nbytes", + "_remote_uniform_nbytes", + "_reuse_tokens", + "_submitting", + "_works", + ) + + def __init__( + self, + endpoint: Endpoint, + handle: object, + local: tuple[Region, ...], + remote: tuple[Region, ...], + indexed: bool, + *, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("TransferPlan objects are created by Endpoint.prepare") + self._endpoint = endpoint + self._handle = handle + self._local = local + self._remote = remote + self._indexed = indexed + self._local_block_count = _logical_block_count(local) + self._remote_block_count = _logical_block_count(remote) + self._local_starts = _catalog_starts(local) + self._remote_starts = _catalog_starts(remote) + self._local_total_nbytes = sum(region.total_nbytes for region in local) + self._remote_total_nbytes = sum(region.total_nbytes for region in remote) + self._local_uniform_nbytes = _uniform_region_nbytes(local) + self._remote_uniform_nbytes = _uniform_region_nbytes(remote) + self._local_uniform_device = _uniform_region_device(local) + self._local_full_nonoverlap = _regions_are_nonoverlapping(local) + self._remote_full_nonoverlap = _regions_are_nonoverlapping(remote) + self._origins = frozenset(region._origin for region in (*local, *remote)) + self._bindings: set[BoundTransfer] = set() + self._reuse_tokens = { + TransferOp.READ: _ReuseToken(), + TransferOp.WRITE: _ReuseToken(), + } + self._submitting: set[object] = set() + self._works: set[Work] = set() + self._creation_failed = False + self._closed = False + + @property + def indexed(self) -> bool: + return self._indexed + + @property + def local_block_count(self) -> int: + """Number of logical blocks in the local indexed catalog.""" + + return self._local_block_count + + @property + def remote_block_count(self) -> int: + """Number of logical blocks in the remote indexed catalog.""" + + return self._remote_block_count + + @property + def closed(self) -> bool: + return self._closed + + def bind( + self, + op: TransferOp | str, + local_indices: Sequence[int] | None = None, + remote_indices: Sequence[int] | None = None, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + ) -> BoundTransfer: + """Validate and retain one reusable submission shape. + + A binding amortizes operation, selection, notification, and CUDA-token + validation across repeated submissions. Each :meth:`BoundTransfer.submit` + still creates an independently owned :class:`Work`. + """ + + return self._endpoint._bind_transfer( + self, + op, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + ) + + def select_indices( + self, + op: TransferOp | str, + *, + local_indices: Sequence[int], + remote_indices: Sequence[int], + ) -> IndexSelection: + """Snapshot and validate one reusable dynamic indexed selection. + + C-contiguous native-int32 buffers use the packed fast input. Mutable + inputs are copied into one private Core-owned buffer, so later caller + mutation cannot change this selection. General sequences remain + supported through ``operator.index`` normalization. + """ + + return self._endpoint._select_plan_indices( + self, + op, + local_indices=local_indices, + remote_indices=remote_indices, + ) + + def submit_indices( + self, + op: TransferOp | str, + *, + local_indices: Sequence[int], + remote_indices: Sequence[int], + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + ) -> Work: + """Snapshot, validate, and submit one dynamic indexed transfer. + + This is the allocation-minimal one-shot form. It creates only Core's + private authenticated backend selection and retains that snapshot in + the returned Work. Use :meth:`select_indices` when the same validated + selection will be bound or reused. + """ + + return self._endpoint._submit_indexed_selection( + plan=self, + op=op, + local_indices=local_indices, + remote_indices=remote_indices, + selection=None, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + + def submit_indices_and_test( + self, + op: TransferOp | str, + *, + local_indices: Sequence[int], + remote_indices: Sequence[int], + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns: int | None = None, + ) -> tuple[Work, bool]: + """Compatibility wrapper returning whether the first state is terminal. + + The Boolean means terminal, including failure or cancellation. A zero + poll budget is post-only. Prefer :meth:`submit_indices_and_poll` when + the caller also needs the exact state; the wrapper does not issue an + additional provider observation. + """ + + work, state = self.submit_indices_and_poll( + op, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + return work, state in _TERMINAL_STATES + + def submit_indices_and_poll( + self, + op: TransferOp | str, + *, + local_indices: Sequence[int], + remote_indices: Sequence[int], + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns: int | None = None, + ) -> tuple[Work, WorkState]: + """Submit one dynamic selection and return its exact bounded-poll state. + + A zero poll budget is post-only. Fused providers return the state from + the post boundary; other providers conservatively return ``PENDING`` + until at least one poll is requested. The returned Work remains the + authoritative owner in every state. + """ + + _validate_bounded_poll_arguments( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=None, + allow_zero=True, + ) + fused = self._endpoint.capabilities.fused_prevalidated_submit_poll + work = self._endpoint._submit_indexed_selection( + plan=self, + op=op, + local_indices=local_indices, + remote_indices=remote_indices, + selection=None, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + initial_max_polls=max_polls if fused else None, + initial_timeout_ns=timeout_ns if fused else None, + ) + if fused: + return work, work._state + if max_polls == 0: + return work, WorkState.PENDING + return work, work.poll_bounded(max_polls=max_polls, timeout_ns=timeout_ns) + + def close(self) -> None: + self._endpoint._close_plan(self) + + def _check_open(self) -> None: + if self._closed: + raise ClosedError("transfer plan is closed") + if self._creation_failed: + raise BackendFailureError("transfer plan is retained for cleanup only") + self._endpoint._check_open() + + def __enter__(self) -> Self: + self._check_open() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class IndexSelection: + """Immutable, exact-plan indexed selection with amortized validation. + + A selection is not an independently closeable resource: its exact plan may + be closed, which invalidates later ``submit`` and ``bind`` calls. Open Work + and BoundTransfer objects retain the selection (and therefore its packed + index snapshot) for as long as a provider may access it. + """ + + __slots__ = ( + "_backend_selection", + "_endpoint", + "_local_devices", + "_op", + "_plan", + "_reuse_token", + "_sealed", + ) + + def __init__( + self, + endpoint: Endpoint, + plan: TransferPlan, + op: TransferOp, + backend_selection: BackendIndexSelection, + local_devices: tuple[str, ...], + *, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("IndexSelection objects are created by TransferPlan") + set_attribute = object.__setattr__ + set_attribute(self, "_endpoint", endpoint) + set_attribute(self, "_plan", plan) + set_attribute(self, "_op", op) + set_attribute(self, "_backend_selection", backend_selection) + set_attribute(self, "_local_devices", local_devices) + # Allocate the provider-neutral plan/op/index-geometry token only if a + # selection-backed BoundTransfer is actually submitted. A one-shot + # selection and a binding used only to create a persistent slot do not + # need a request-pool identity. + set_attribute(self, "_reuse_token", None) + set_attribute(self, "_sealed", True) + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError("IndexSelection is immutable") + + def __delattr__(self, name: str) -> None: + raise AttributeError("IndexSelection is immutable") + + @property + def op(self) -> TransferOp: + return self._op + + @property + def local_count(self) -> int: + return len(self._backend_selection.local_indices) + + @property + def remote_count(self) -> int: + return len(self._backend_selection.remote_indices) + + @property + def selected_nbytes(self) -> int: + return self._backend_selection.selected_nbytes + + def submit( + self, + *, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + ) -> Work: + """Submit this selection with request-local completion arguments.""" + + return self._endpoint._submit_indexed_selection( + plan=None, + op=None, + local_indices=None, + remote_indices=None, + selection=self, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + + def submit_and_test( + self, + *, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns: int | None = None, + ) -> tuple[Work, bool]: + """Compatibility wrapper returning whether the first state is terminal. + + The Boolean has the same meaning as :meth:`Work.test`: ``True`` means + terminal, including failure or cancellation, rather than successful. + Prefer :meth:`submit_and_poll` when the exact state is needed. + """ + + work, state = self.submit_and_poll( + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + return work, state in _TERMINAL_STATES + + def submit_and_poll( + self, + *, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns: int | None = None, + ) -> tuple[Work, WorkState]: + """Submit and return the exact state from a bounded initial poll. + + A zero poll budget is post-only. On a provider without fused support, + Core conservatively returns ``PENDING`` because no post-derived state + has been observed. The returned Work remains authoritative. + """ + + _validate_bounded_poll_arguments( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=None, + allow_zero=True, + ) + if self._endpoint.capabilities.fused_prevalidated_submit_poll: + work = self._endpoint._submit_indexed_selection( + plan=None, + op=None, + local_indices=None, + remote_indices=None, + selection=self, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + initial_max_polls=max_polls, + initial_timeout_ns=timeout_ns, + ) + return work, work._state + + work = self.submit( + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + if max_polls == 0: + return work, WorkState.PENDING + return work, work.poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + + def bind( + self, + *, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + ) -> BoundTransfer: + """Bind completion arguments without revalidating this selection.""" + + return self._endpoint._bind_index_selection( + self, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + ) + + def _check_open(self) -> TransferPlan: + if self._plan._closed: + raise ClosedError("transfer plan is closed") + self._plan._check_open() + return self._plan + + +class BoundTransfer: + """A prevalidated, reusable transfer submission. + + Bindings retain their plan and canonical submission arguments. They are + Core-only resources: providers observe only the binding's opaque, stable + reuse token on submission. + """ + + __slots__ = ( + "_closed", + "_cuda_device", + "_cuda_ordering", + "_endpoint", + "_local_indices", + "_notification", + "_op", + "_plan", + "_ready_event", + "_remote_indices", + "_reuse_token", + "_selection", + "_slots", + "_stream", + "_submitting", + "_works", + ) + + def __init__( + self, + endpoint: Endpoint, + plan: TransferPlan, + op: TransferOp, + *, + local_indices: tuple[int, ...] | None, + remote_indices: tuple[int, ...] | None, + notification: bytes | None, + cuda_device: torch.device | None, + cuda_ordering: CudaOrderingMode, + stream: object | None, + ready_event: object | None, + selection: IndexSelection | None = None, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("BoundTransfer objects are created by TransferPlan.bind") + self._endpoint = endpoint + self._plan: TransferPlan | None = plan + self._op = op + self._local_indices = local_indices + self._remote_indices = remote_indices + self._notification = notification + self._cuda_device = cuda_device + self._cuda_ordering = cuda_ordering + self._stream = stream + self._ready_event = ready_event + self._selection = selection + # Selection-backed bindings share the selection's lazily allocated + # geometry token. Do not allocate a redundant binding token for them. + self._reuse_token: object | None = _ReuseToken() if selection is None else None + self._submitting: set[Work] = set() + self._works: set[Work] = set() + self._slots: set[TransferSlot] = set() + self._closed = False + + @property + def op(self) -> TransferOp: + return self._op + + @property + def closed(self) -> bool: + return self._closed + + def submit( + self, + *, + notification: bytes | bytearray | memoryview | None = cast( + Any, _UNSET_NOTIFICATION + ), + recovery_token: object | None = None, + ) -> Work: + """Submit the binding. + + Selection-backed bindings accept a request-local notification override + so one native indexed lowering can serve successive framework tags. + Ordinary bindings retain their original fully immutable contract. + """ + + return self._endpoint._submit_bound( + self, + notification=notification, + recovery_token=recovery_token, + ) + + def create_slot( + self, *, request_local_notifications: bool = False + ) -> TransferSlot | RequestLocalNotificationSlot: + """Create a persistent, single-flight request for this binding. + + The default returns the fixed-notification :class:`TransferSlot` and + preserves its minimal execute path. ``request_local_notifications=True`` + returns a :class:`RequestLocalNotificationSlot`, whose start operation + requires an explicit payload (or ``None`` to clear it) for every + generation. This opt-in mode is capability gated and requires a + binding with no fixed notification. + """ + + if not isinstance(request_local_notifications, bool): + raise TypeError("request_local_notifications must be a bool") + return self._endpoint._create_transfer_slot( + self, + request_local_notifications=request_local_notifications, + ) + + def close(self) -> None: + self._endpoint._close_bound_transfer(self) + + def _check_open(self) -> TransferPlan: + if self._closed: + raise ClosedError("bound transfer is closed") + plan = self._plan + if plan is None: + raise ClosedError("bound transfer is closed") + plan._check_open() + return plan + + def __enter__(self) -> Self: + self._check_open() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class TransferSlot: + """Stable, mutable single-flight request for one ``BoundTransfer``. + + A slot owns at most one unrecycled generation. ``state`` is ``None`` while + idle and a :class:`WorkState` after ``start`` until terminal ``recycle``. + ``active`` deliberately remains true for a terminal generation: recycling + is the explicit boundary that permits the next start and prevents ABA-style + observations across generations. ``execute`` is the opt-in full-cycle form; + it publishes a clean successful generation directly back to idle. + """ + + __slots__ = ( + "_active", + "_backend_cancel", + "_backend_close", + "_backend_poll_bounded", + "_backend_poll_state", + "_backend_recycle", + "_backend_deferred_successful_recycle", + "_backend_execution_dispatch", + "_backend_generation_active", + "_caller_ready", + "_ready_event_synchronize", + "_fast_execute", + "_backend_slot", + "_backend_start", + "_backend_start_and_poll", + "_bound", + "_closed", + "_creation_failed", + "_endpoint", + "_error", + "_generation", + "_lock", + "_owner_thread", + "_start_error", + "_state", + ) + + def __init__( + self, + endpoint: Endpoint, + bound: BoundTransfer, + *, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("TransferSlot objects are created by BoundTransfer") + self._endpoint = endpoint + self._bound: BoundTransfer | None = bound + self._backend_slot: BackendRequestSlot | None = None + self._backend_execution_dispatch: object | None = None + self._backend_start: Callable[[], None] | None = None + self._backend_start_and_poll: Callable[..., WorkState | None] | None = None + self._backend_poll_bounded: Callable[..., WorkState | None] | None = None + self._backend_poll_state: Callable[[], WorkState | None] | None = None + self._backend_cancel: Callable[[], bool] | None = None + self._backend_recycle: Callable[[], None] | None = None + self._backend_close: Callable[[], None] | None = None + self._backend_deferred_successful_recycle = ( + endpoint.capabilities.deferred_successful_request_slot_recycle + ) + self._backend_generation_active = False + self._caller_ready = bound._cuda_ordering is CudaOrderingMode.CALLER_READY + # A BoundTransfer has already authenticated the token type and device. + # Pin one complete readiness operation per persistent slot. An explicit + # event contributes its bound synchronize method; a stream contributes + # one reusable, lazily materialized CUDA event plus cached capture, + # record, and synchronize callables. Every generation still performs + # capture detection, records after prior producer work, and host-waits, + # but it allocates neither a Python nor native CUDA event in steady state. + if self._caller_ready: + self._ready_event_synchronize = None + elif bound._stream is not None: + self._ready_event_synchronize = _make_reusable_cuda_stream_readiness( + bound._stream, + bound._cuda_device, + ) + else: + self._ready_event_synchronize = cast( + Callable[[], None], bound._ready_event.synchronize + ) + self._fast_execute: Callable[..., WorkState | None] | None = None + # SINGLE and CALLER_SERIALIZED use a genuinely lock-free path. ``None`` + # avoids even the Python context-manager calls paid by Core's general + # no-op lock. The latter contract requires whole-graph exclusion by its + # caller; it is not an independently thread-safe slot. + self._lock = ( + None + if endpoint._thread_mode + in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) + else threading.RLock() + ) + self._owner_thread = ( + endpoint._owner_thread + if endpoint._thread_mode is ThreadMode.SINGLE + else None + ) + self._generation = 0 + self._active = False + self._state: WorkState | None = None + self._error: BaseException | None = None + self._start_error: BaseException | None = None + self._creation_failed = False + self._closed = False + + @property + def generation(self) -> int: + """Number of starts owned by this slot, including retained failures.""" + + return self._generation + + @property + def active(self) -> bool: + """Whether this slot owns a generation that has not been recycled.""" + + return self._active + + @property + def closed(self) -> bool: + return self._closed + + @property + def state(self) -> WorkState | None: + return self._refresh() + + @property + def error(self) -> BaseException | None: + lock = self._lock + if lock is None: + self._check_access() + state = self._refresh_unlocked() + return None if state is None else self._error + with lock: + self._check_access() + state = self._refresh_unlocked() + return None if state is None else self._error + + @property + def completion_event(self) -> object | None: + lock = self._lock + if lock is None: + return self._completion_event_unlocked() + with lock: + return self._completion_event_unlocked() + + def start(self) -> None: + """Start one generation after re-establishing CUDA readiness. + + Once Core enters the provider ``start`` call, the generation is owned + even if an exception crosses that boundary. Such an exception is + retained and becomes visible as terminal failure; the slot cannot be + reused or closed until its provider state is terminal and recycled. + """ + + lock = self._lock + if lock is None: + self._start_unlocked() + return + with lock: + self._start_unlocked() + + def start_and_test(self) -> bool: + """Start and logically recycle a post-reported immediate success. + + ``True`` means the generation completed successfully and the public + slot is already idle/repostable. ``False`` means it remains active and + can be observed with :meth:`test_and_recycle`. Terminal failure is + raised and deliberately remains active for explicit error recovery. + """ + + fast_start_and_test = self._fast_execute + if fast_start_and_test is not None: + next_generation = self._generation + 1 + try: + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError( + "transfer slot is retained for cleanup only" + ) + if ( + self._owner_thread is not None + and _GET_IDENT() != self._owner_thread + ): + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + if self._active: + raise BusyError( + "transfer slot already owns an unrecycled generation" + ) + + if not self._caller_ready: + ready_event_synchronize = self._ready_event_synchronize + if ready_event_synchronize is not None: + try: + ready_event_synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to wait for CUDA readiness" + ) from error + else: + bound = self._bound + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + try: + self._generation = next_generation + self._active = True + except BaseException: + self._active = False + self._generation = next_generation - 1 + self._state = None + self._error = None + self._start_error = None + self._backend_generation_active = False + raise + + state = fast_start_and_test(max_polls=0, timeout_ns=None) + except BaseException as error: + if self._generation != next_generation: + raise + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + if state is _WORK_COMPLETED: + self._active = False + return True + if state is _WORK_RUNNING: + # The fused provider has acknowledged this generation and + # reported the dominant QD-many nonterminal state. Avoid + # reclassifying the same enum through the generic observer; + # these two stores are exactly its RUNNING-state commit. + self._backend_generation_active = True + self._state = state + return False + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot start returned without " + "acknowledging a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + elif not isinstance(state, WorkState): + self._start_error = BackendFailureError( + "backend returned invalid fused request-slot state " + f"{state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + state = self._observe_backend_state_unlocked(state) + + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + start_and_poll = self._backend_start_and_poll + lock = self._lock + if start_and_poll is None: + if lock is None: + return self._start_and_test_fallback_unlocked() + with lock: + return self._start_and_test_fallback_unlocked() + if lock is None: + state = self._start_and_poll_unlocked(start_and_poll, None, max_polls=0) + expected_generation = self._generation + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + with lock: + state = self._start_and_poll_unlocked(start_and_poll, None, max_polls=0) + expected_generation = self._generation + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def _start_and_test_fallback_unlocked(self) -> bool: + """Keep fallback start/observe/finalize in one slot transaction.""" + + self._start_unlocked() + expected_generation = self._generation + state = self._refresh_unlocked() + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def test_and_recycle(self) -> bool: + """Observe once and logically recycle only successful completion. + + Returns ``False`` while the generation remains pending/running. + Successful completion returns ``True`` with the public slot idle. + Failure or cancellation is raised and remains active for explicit + recovery. Deferred-success providers retain the clean native receipt + for next-start/close without a provider recycle call. + """ + + lock = self._lock + if lock is None and self._backend_deferred_successful_recycle: + # Exact lock-free/deferred continuation path. Slot construction has + # already authenticated the lifetime graph and cached the provider + # dispatch. Keep the common PENDING/RUNNING/COMPLETED states on one + # straight-line poll without the generic refresh/observe/finalize + # frames; failure and malformed-provider states retain the full + # diagnostic path below. + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError("transfer slot is retained for cleanup only") + if self._owner_thread is not None and _GET_IDENT() != self._owner_thread: + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + if not self._active: + raise BusyError("transfer slot has no active generation") + if self._state is WorkState.COMPLETED: + # A preceding state/test observation may already have cached + # successful completion. Commit Core-idle directly; polling + # the provider again cannot add information and penalizes the + # normal test-then-recycle composition. + self._fast_execute = None + self._backend_generation_active = False + self._error = None + self._start_error = None + self._state = None + # Deferred success leaves a durable provider receipt, so all + # private Core state can be clean before the public idle commit. + # This prevents an asynchronous exception from exposing stale + # terminal state beside a subclass-specific fast dispatcher. + self._active = False + self._restore_fast_execute_unlocked() + return True + if self._state in _TERMINAL_STATES: + self._raise_for_state() + poll_state = self._backend_poll_state + if poll_state is None: + raise BackendFailureError( + "transfer slot has no backend state-poll operation" + ) + try: + state = poll_state() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._endpoint.backend!r} failed to query " + "persistent request state" + ) from error + if state is WorkState.COMPLETED and self._start_error is None: + self._state = WorkState.COMPLETED + self._backend_generation_active = True + self._fast_execute = None + self._backend_generation_active = False + self._error = None + self._start_error = None + self._state = None + # Keep the deferred-success idle commit last for the same + # interruption-safety reason as the cached-completion branch. + self._active = False + self._restore_fast_execute_unlocked() + return True + if state is None: + if self._start_error is None: + self._start_error = BackendFailureError( + "backend request-slot start returned without " + "acknowledging a generation" + ) + self._error = self._start_error + self._state = WorkState.FAILED + self._raise_for_state() + if not isinstance(state, WorkState): + raise BackendFailureError( + f"backend returned invalid transfer slot state {state!r}" + ) + if self._state is WorkState.RUNNING and state is WorkState.PENDING: + raise BackendFailureError( + "backend transfer slot state regressed from running to pending" + ) + self._backend_generation_active = True + if state is WorkState.PENDING or state is WorkState.RUNNING: + self._state = state + return False + self._observe_backend_state_unlocked(state) + self._raise_for_state() + + if lock is None: + self._check_access() + if not self._active: + raise BusyError("transfer slot has no active generation") + expected_generation = self._generation + state = self._refresh_unlocked() + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + with lock: + self._check_access() + if not self._active: + raise BusyError("transfer slot has no active generation") + expected_generation = self._generation + state = self._refresh_unlocked() + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def test(self) -> bool: + state = self._refresh() + return state in _TERMINAL_STATES if state is not None else False + + def wait(self, timeout: float | None = None) -> TransferSlot: + deadline_ns = _transfer_slot_deadline_ns(timeout) + lock = self._lock + if lock is None: + expected_generation = self._generation_for_wait_unlocked() + else: + with lock: + expected_generation = self._generation_for_wait_unlocked() + return self._wait_until( + deadline_ns, + spin=False, + expected_generation=expected_generation, + ) + + def start_and_wait(self, timeout: float | None = None) -> TransferSlot: + """Start one generation and wait, using bounded provider fusion. + + The optional timeout is a total host-side budget beginning at method + entry. It cannot preempt a blocking CUDA-readiness or provider call. + Timeout leaves the generation active so it can be observed, cancelled, + and explicitly recycled after reaching terminal state. + """ + + deadline_ns = _transfer_slot_deadline_ns(timeout) + start_and_poll = self._backend_start_and_poll + poll_bounded = self._backend_poll_bounded + if start_and_poll is None: + lock = self._lock + if lock is None: + self._start_unlocked() + expected_generation = self._generation + else: + with lock: + self._start_unlocked() + expected_generation = self._generation + if poll_bounded is not None: + return self._wait_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_until( + deadline_ns, spin=True, expected_generation=expected_generation + ) + + lock = self._lock + if lock is None: + state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self + if state in _TERMINAL_STATES: + return self._raise_for_state() + else: + with lock: + state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self + if state in _TERMINAL_STATES: + return self._raise_for_state() + + # A provider is allowed to exhaust its small bounded poll budget while + # work remains active. Do not retain the per-slot lock across the + # potentially long ordinary wait. + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + # Preserve the zero-timeout no-extra-poll contract without letting + # a concurrent recycle/restart be misreported as a timeout. + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if poll_bounded is not None: + return self._wait_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_until( + deadline_ns, + spin=True, + expected_generation=expected_generation, + ) + + def execute(self, timeout: float | None = None) -> int: + """Start, wait, and logically recycle one successful generation. + + The returned integer identifies the completed generation. Success makes + the public slot idle before returning, so ``state`` is ``None`` and a + later generation may start immediately. Failure, cancellation, and + timeout deliberately retain the active generation for observation, + cancellation, and explicit recycle. + + Providers advertising deferred successful recycle may fold their + successful provider-side recycle into the next start or close boundary. + Other providers are recycled before Core publishes the slot as idle. + """ + + fast_execute = self._fast_execute + + # Specialize the dominant small-transfer lifecycle as one straight-line + # Core transaction: lock-free ownership, an unbounded fused provider + # call, and a provider that retains clean successful receipts. This + # preserves the generic transaction's checks and interruption + # boundaries while avoiding four Python helper frames on immediate + # completion. + if timeout is None and fast_execute is not None: + next_generation = self._generation + 1 + try: + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError( + "transfer slot is retained for cleanup only" + ) + if ( + self._owner_thread is not None + and _GET_IDENT() != self._owner_thread + ): + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + if self._active: + raise BusyError( + "transfer slot already owns an unrecycled generation" + ) + + if not self._caller_ready: + ready_event_synchronize = self._ready_event_synchronize + if ready_event_synchronize is not None: + try: + ready_event_synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to wait for CUDA readiness" + ) from error + else: + bound = self._bound + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + try: + # Idle already guarantees clear diagnostics and no active + # provider generation. `_active` is the sole publication + # commit; the provider dispatch owns durable request + # state after entry, so this synchronous SINGLE path needs + # no transient Core PENDING store. + self._generation = next_generation + self._active = True + except BaseException: + self._active = False + self._generation = next_generation - 1 + self._state = None + self._error = None + self._start_error = None + self._backend_generation_active = False + raise + + state = fast_execute( + max_polls=_FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns=None, + ) + except BaseException as error: + if self._generation != next_generation: + raise + # Once Core publishes the generation, a provider call may have + # launched native work. Preserve the same ambiguous-start + # recovery used by the generic fused path. + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + if state is _WORK_COMPLETED: + # The provider completion receipt is already durable. Before + # idle commit, an active Core generation rediscovers it; + # after commit, Core is already clean-idle. No transient + # Core state stores are needed on this exact path. + self._active = False + return next_generation + if state is _WORK_RUNNING: + # The exact RUNNING singleton is already a valid provider + # acknowledgement. Publish provider ownership before the + # observable state, then enter the cached continuation + # directly instead of reclassifying the initial result. + self._backend_generation_active = True + self._state = state + poll_bounded = self._backend_poll_bounded + if poll_bounded is not None: + return self._wait_execute_bounded_until( + None, + expected_generation=next_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + None, expected_generation=next_generation + ) + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot start returned without " + "acknowledging a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + elif not isinstance(state, WorkState): + self._start_error = BackendFailureError( + "backend returned invalid fused request-slot state " + f"{state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + state = self._observe_backend_state_unlocked(state) + + expected_generation = self._generation + if state in _TERMINAL_STATES: + self._raise_for_state() + poll_bounded = self._backend_poll_bounded + if poll_bounded is not None: + return self._wait_execute_bounded_until( + None, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + None, expected_generation=expected_generation + ) + + # Non-fast paths preserve the shared finite-timeout validator. Keeping + # this below the fast return avoids a deadline local on the dominant + # unbounded deferred SINGLE completion. + deadline_ns = None if timeout is None else _transfer_slot_deadline_ns(timeout) + start_and_poll = self._backend_start_and_poll + lock = self._lock + poll_bounded = self._backend_poll_bounded + if start_and_poll is None: + if lock is None: + self._start_unlocked() + expected_generation = self._generation + else: + with lock: + self._start_unlocked() + expected_generation = self._generation + if poll_bounded is not None: + return self._wait_execute_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + deadline_ns, expected_generation=expected_generation + ) + + if lock is None: + state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) + if state is WorkState.COMPLETED: + return self._finish_execute_success_unlocked(self._generation) + expected_generation = self._generation + if state in _TERMINAL_STATES: + self._raise_for_state() + else: + with lock: + state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self._finish_execute_success_unlocked(expected_generation) + if state in _TERMINAL_STATES: + self._raise_for_state() + + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if poll_bounded is not None: + return self._wait_execute_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + deadline_ns, expected_generation=expected_generation + ) + + def _wait_until( + self, + deadline_ns: int | None, + *, + spin: bool, + expected_generation: int, + ) -> TransferSlot: + polls = 0 + while True: + if self._observe_wait_generation(expected_generation): + return self + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + raise TransferTimeoutError("transfer slot timed out") + if self._endpoint._progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + polls += 1 + if not spin and polls >= 16: + time.sleep(0 if polls < 128 else 0.0001) + + def _wait_execute_until( + self, + deadline_ns: int | None, + *, + expected_generation: int, + ) -> int: + """Spin until success and publish Core-idle in the observation lock.""" + + while True: + lock = self._lock + if lock is None: + completed_generation = self._observe_execute_generation_unlocked( + expected_generation + ) + else: + with lock: + completed_generation = self._observe_execute_generation_unlocked( + expected_generation + ) + if completed_generation is not None: + return completed_generation + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if self._endpoint._progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + + def _wait_bounded_until( + self, + deadline_ns: int | None, + *, + expected_generation: int, + poll_bounded: Callable[..., WorkState | None], + ) -> TransferSlot: + lock = self._lock + while True: + if lock is None: + terminal = self._poll_bounded_generation_unlocked( + expected_generation, deadline_ns, poll_bounded + ) + else: + with lock: + terminal = self._poll_bounded_generation_unlocked( + expected_generation, deadline_ns, poll_bounded + ) + if terminal: + return self + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + raise TransferTimeoutError("transfer slot timed out") + if self._endpoint._progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + + def _wait_execute_bounded_until( + self, + deadline_ns: int | None, + *, + expected_generation: int, + poll_bounded: Callable[..., WorkState | None], + ) -> int: + """Run bounded polls and publish successful recycle under the slot lock.""" + + lock = self._lock + while True: + if lock is None: + terminal = self._poll_bounded_generation_unlocked( + expected_generation, deadline_ns, poll_bounded + ) + if terminal: + return self._finish_execute_success_unlocked(expected_generation) + else: + with lock: + terminal = self._poll_bounded_generation_unlocked( + expected_generation, deadline_ns, poll_bounded + ) + if terminal: + return self._finish_execute_success_unlocked( + expected_generation + ) + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if self._endpoint._progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + + def _poll_bounded_generation_unlocked( + self, + expected_generation: int, + deadline_ns: int | None, + poll_bounded: Callable[..., WorkState | None], + ) -> bool: + self._validate_wait_generation_unlocked(expected_generation) + if self._state is WorkState.COMPLETED: + return True + if self._state in _TERMINAL_STATES: + self._raise_for_state() + configured_timeout_ns = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS + if deadline_ns is None: + timeout_ns = configured_timeout_ns + else: + remaining_ns = max(0, deadline_ns - time.monotonic_ns()) + timeout_ns = ( + remaining_ns + if configured_timeout_ns is None + else min(configured_timeout_ns, remaining_ns) + ) + if self._lock is None: + try: + state = poll_bounded( + max_polls=_FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + timeout_ns=timeout_ns, + timeout_check_interval=( + _FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL + ), + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._endpoint.backend!r} failed during bounded " + "persistent request polling" + ) from error + else: + state = self._endpoint._backend_request_slot_poll_bounded( + poll_bounded, + max_polls=_FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + timeout_ns=timeout_ns, + timeout_check_interval=( + _FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL + ), + ) + state = self._observe_backend_state_unlocked(state) + if state is WorkState.COMPLETED: + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def cancel(self) -> bool: + lock = self._lock + if lock is None: + return self._cancel_unlocked() + with lock: + return self._cancel_unlocked() + + def recycle(self) -> None: + """Return a terminal generation to idle without changing its shape.""" + + lock = self._lock + if lock is None: + self._recycle_unlocked() + return + with lock: + self._recycle_unlocked() + + def close(self) -> None: + lock = self._lock + if lock is None: + self._endpoint._close_transfer_slot(self) + return + with lock: + self._endpoint._close_transfer_slot(self) + + def _adopt_backend_slot( + self, + backend_slot: BackendRequestSlot, + *, + execution_dispatch: object | None = None, + ) -> None: + if not isinstance(backend_slot, BackendRequestSlot): + raise BackendFailureError( + "backend must adopt a BackendRequestSlot instance" + ) + lock = self._lock + if lock is None: + self._adopt_backend_slot_unlocked( + backend_slot, execution_dispatch=execution_dispatch + ) + return + with lock: + self._adopt_backend_slot_unlocked( + backend_slot, execution_dispatch=execution_dispatch + ) + + def _refresh(self) -> WorkState | None: + lock = self._lock + if lock is None: + self._check_access() + return self._refresh_unlocked() + with lock: + self._check_access() + return self._refresh_unlocked() + + def _start_unlocked(self) -> None: + previous_generation = self._generation + try: + self._begin_generation_unlocked() + start = self._backend_start + if start is None: + self._start_error = BackendFailureError( + "transfer slot has no adopted backend request" + ) + return + if self._lock is None: + result = start() + else: + result = self._endpoint._backend_request_slot_start(start) + # A normal return acknowledges that provider start was entered, + # even when the provider then violates the required None return. + self._backend_generation_active = True + if result is not None: + raise BackendFailureError("backend request-slot start must return None") + except BaseException as error: + if self._generation == previous_generation: + # CUDA readiness or an interrupted pre-commit publication did + # not create a generation and must remain a direct exception. + raise + # Native launch is ambiguous once the provider call has begun. + # Preserve ownership; a conforming provider exposes terminal state + # or permits cancellation rather than stranding this generation. + self._start_error = error + + def _begin_generation_unlocked(self) -> None: + bound = self._check_access() + if self._active: + raise BusyError("transfer slot already owns an unrecycled generation") + + if self._state is not None: + self._backend_generation_active = False + self._error = None + self._start_error = None + self._state = None + + # CUDA readiness failure is unambiguously before provider activity, so + # it does not consume a generation. This runs outside endpoint state + # and provider locks and is repeated for every start. + if not self._caller_ready: + ready_event_synchronize = self._ready_event_synchronize + if ready_event_synchronize is not None: + try: + ready_event_synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to wait for CUDA readiness" + ) from error + else: + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + previous_generation = self._generation + try: + # A generic/mixed generation temporarily disables the exact execute + # path. Its idle transition restores the cached callable only after + # all private state is clean. CUDA readiness precedes this store, so + # an unambiguous readiness failure preserves fast eligibility. + self._fast_execute = None + # ``_active`` is the publication commit and is deliberately last. + # An interrupt before it leaves an idle slot; an interrupt after it + # leaves a complete, self-consistent Core generation that the + # caller can observe and recycle without touching an idle provider. + self._state = WorkState.PENDING + self._error = None + self._start_error = None + self._backend_generation_active = False + self._generation = previous_generation + 1 + self._active = True + except BaseException: + self._active = False + self._generation = previous_generation + self._state = None + self._error = None + self._start_error = None + self._backend_generation_active = False + self._restore_fast_execute_unlocked() + raise + + def _start_and_poll_unlocked( + self, + start_and_poll: Callable[..., WorkState | None], + deadline_ns: int | None, + *, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + ) -> WorkState: + previous_generation = self._generation + try: + self._begin_generation_unlocked() + timeout_ns = ( + None + if deadline_ns is None + else min( + _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS, + max(0, deadline_ns - time.monotonic_ns()), + ) + ) + if self._lock is None: + state = start_and_poll( + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + else: + state = self._endpoint._backend_request_slot_start_and_poll( + start_and_poll, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + except BaseException as error: + if self._generation == previous_generation: + raise + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot start returned without acknowledging " + "a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + if not isinstance(state, WorkState): + self._start_error = BackendFailureError( + f"backend returned invalid fused request-slot state {state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + return self._observe_backend_state_unlocked(state) + + def _completion_event_unlocked(self) -> object | None: + self._check_access() + if not self._active: + raise BusyError("transfer slot has no active generation") + if not self._endpoint.capabilities.cuda_completion_event: + return None + backend_slot = self._require_backend_slot() + return self._endpoint._backend_request_slot_completion_event(backend_slot) + + def _generation_for_wait_unlocked(self) -> int: + self._check_access() + if not self._active: + raise BusyError("transfer slot has no active generation") + return self._generation + + def _observe_wait_generation(self, expected_generation: int) -> bool: + lock = self._lock + if lock is None: + return self._observe_wait_generation_unlocked(expected_generation) + with lock: + return self._observe_wait_generation_unlocked(expected_generation) + + def _validate_wait_generation(self, expected_generation: int) -> None: + lock = self._lock + if lock is None: + self._validate_wait_generation_unlocked(expected_generation) + return + with lock: + self._validate_wait_generation_unlocked(expected_generation) + + def _validate_wait_generation_unlocked(self, expected_generation: int) -> None: + self._check_access() + if not self._active or self._generation != expected_generation: + raise BusyError( + f"transfer slot generation {expected_generation} changed while " + "waiting" + ) + + def _observe_wait_generation_unlocked(self, expected_generation: int) -> bool: + self._check_access() + if not self._active or self._generation != expected_generation: + raise BusyError( + f"transfer slot generation {expected_generation} changed while " + "waiting" + ) + state = self._refresh_unlocked() + if state in _TERMINAL_STATES: + self._raise_for_state() + return True + return False + + def _observe_execute_generation_unlocked( + self, expected_generation: int + ) -> int | None: + self._validate_wait_generation_unlocked(expected_generation) + state = self._refresh_unlocked() + if state is WorkState.COMPLETED: + return self._finish_execute_success_unlocked(expected_generation) + if state in _TERMINAL_STATES: + self._raise_for_state() + return None + + def _cancel_unlocked(self) -> bool: + self._check_access() + if not self._active: + return False + if self._refresh_unlocked() in _TERMINAL_STATES: + return self._state is WorkState.CANCELLED + cancel = self._backend_cancel + if cancel is None: + raise BackendFailureError("transfer slot has no adopted backend request") + cancelled = self._endpoint._backend_request_slot_cancel(cancel) + if not isinstance(cancelled, bool): + raise BackendFailureError("backend slot cancellation must return a bool") + self._refresh_unlocked() + return cancelled and self._state is WorkState.CANCELLED + + def _recycle_unlocked(self) -> None: + self._check_access() + if not self._active: + raise BusyError("transfer slot has no generation to recycle") + self._fast_execute = None + if self._refresh_unlocked() not in _TERMINAL_STATES: + raise BusyError("cannot recycle a non-terminal transfer slot") + if self._backend_generation_active: + recycle = self._backend_recycle + if recycle is None: + raise BackendFailureError( + "transfer slot has no adopted backend request" + ) + if self._lock is None: + try: + result = recycle() + if result is not None: + raise BackendFailureError( + "backend request-slot recycle must return None" + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._endpoint.backend!r} failed to recycle a " + "persistent request" + ) from error + else: + self._endpoint._backend_request_slot_recycle(recycle) + # ``_active = False`` is the sole public idle-state commit. Private + # terminal diagnostics may remain briefly if an asynchronous exception + # lands after it; public state/error ignore them while idle, and the + # next generation overwrites them before its own active commit. + self._active = False + self._backend_generation_active = False + self._error = None + self._start_error = None + # Clear the dirty-idle marker only after every other private field. + self._state = None + self._restore_fast_execute_unlocked() + + def _finish_execute_success_unlocked(self, expected_generation: int) -> int: + """Commit a successful full-cycle generation while holding its lock.""" + + # Every caller has already validated access and generation while holding + # this slot's serialization boundary. Retain the cheap inline ABA guard + # without repeating `_check_access()` on every successful generation. + if not self._active or self._generation != expected_generation: + raise BusyError( + f"transfer slot generation {expected_generation} changed while " + "waiting" + ) + if self._state is not WorkState.COMPLETED: + raise BackendFailureError( + "transfer slot execute success requires completed provider state" + ) + self._fast_execute = None + if ( + self._backend_generation_active + and not self._backend_deferred_successful_recycle + ): + recycle = self._backend_recycle + if recycle is None: + raise BackendFailureError( + "transfer slot has no adopted backend request" + ) + if self._lock is None: + try: + result = recycle() + if result is not None: + raise BackendFailureError( + "backend request-slot recycle must return None" + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._endpoint.backend!r} failed to recycle a " + "persistent request" + ) from error + else: + self._endpoint._backend_request_slot_recycle(recycle) + + completed_generation = self._generation + # The provider retains an observable COMPLETED receipt when its recycle + # is deferred. ``_active = False`` is the sole public idle commit; the + # private terminal state is cleared afterward. Restore the exact fast + # callable only after every private field is clean. + self._active = False + self._backend_generation_active = False + self._error = None + self._start_error = None + self._state = None + self._restore_fast_execute_unlocked() + return completed_generation + + def _restore_fast_execute_unlocked(self) -> None: + """Publish exact-execute eligibility after a completely clean idle.""" + + start_and_poll = self._backend_start_and_poll + self._fast_execute = ( + start_and_poll + if ( + self._lock is None + and start_and_poll is not None + and self._backend_deferred_successful_recycle + ) + else None + ) + + def _adopt_backend_slot_unlocked( + self, + backend_slot: BackendRequestSlot, + *, + execution_dispatch: object | None, + ) -> None: + if self._closed: + raise BackendFailureError( + "backend attempted to adopt an expired transfer slot" + ) + if self._backend_slot is not None: + raise BackendFailureError( + "backend attempted to adopt transfer slot more than once" + ) + # Allocate every bound method before publishing `_backend_slot` as the + # adoption commit marker. An asynchronous interruption therefore leaves + # either a wholly unadopted slot or a fully usable one, never a public + # slot with only part of its steady-state dispatch table installed. + dispatch = backend_slot if execution_dispatch is None else execution_dispatch + start = dispatch.start + poll_state = dispatch.poll_state + start_and_poll = ( + dispatch.start_and_poll + if self._endpoint.capabilities.fused_request_slot_start_poll + else None + ) + poll_bounded = ( + dispatch.poll_bounded + if self._endpoint.capabilities.fused_request_slot_poll + else None + ) + cancel = dispatch.cancel + recycle = dispatch.recycle + close = backend_slot.close + operations = ( + ("start", start), + ("poll_state", poll_state), + ("cancel", cancel), + ("recycle", recycle), + ("close", close), + ) + if self._endpoint.capabilities.fused_request_slot_start_poll: + operations += (("start_and_poll", start_and_poll),) + if self._endpoint.capabilities.fused_request_slot_poll: + operations += (("poll_bounded", poll_bounded),) + for name, operation in operations: + if not callable(operation): + raise BackendFailureError( + f"backend request-slot {name} operation must be callable" + ) + self._backend_start = start + self._backend_poll_state = poll_state + self._backend_start_and_poll = start_and_poll + self._backend_poll_bounded = poll_bounded + self._backend_cancel = cancel + self._backend_recycle = recycle + self._backend_close = close + self._backend_execution_dispatch = dispatch + self._fast_execute = ( + start_and_poll + if ( + self._lock is None + and start_and_poll is not None + and self._backend_deferred_successful_recycle + ) + else None + ) + self._backend_slot = backend_slot + + def _refresh_unlocked(self) -> WorkState | None: + if not self._active: + return None + if self._state is WorkState.COMPLETED: + return self._state + if self._state in _TERMINAL_STATES: + return self._state + poll_state = self._backend_poll_state + if poll_state is None: + raise BackendFailureError( + "transfer slot has no backend state-poll operation" + ) + if self._lock is None: + try: + state = poll_state() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._endpoint.backend!r} failed to query " + "persistent request state" + ) from error + else: + state = self._endpoint._backend_request_slot_state(poll_state) + if state is None: + if self._start_error is None: + self._start_error = BackendFailureError( + "backend request-slot start returned without acknowledging " + "a generation" + ) + self._error = self._start_error + self._state = WorkState.FAILED + return self._state + return self._observe_backend_state_unlocked(state) + + def _observe_backend_state_unlocked(self, state: object) -> WorkState: + if not isinstance(state, WorkState): + raise BackendFailureError( + f"backend returned invalid transfer slot state {state!r}" + ) + if self._state is WorkState.RUNNING and state is WorkState.PENDING: + raise BackendFailureError( + "backend transfer slot state regressed from running to pending" + ) + if not self._backend_generation_active: + self._backend_generation_active = True + if state is WorkState.COMPLETED and self._start_error is None: + self._state = state + return state + observed_state = state + observed_error: BaseException | None = None + if state in _TERMINAL_STATES and self._start_error is not None: + observed_state = WorkState.FAILED + observed_error = self._start_error + elif state is WorkState.FAILED: + try: + backend_slot = self._require_backend_slot() + error = self._endpoint._backend_request_slot_error(backend_slot) + except Exception as observation_error: + observed_error = observation_error + else: + if error is None: + error = BackendFailureError( + "transfer backend reported slot failure without an error" + ) + elif not isinstance(error, BaseException): + error = BackendFailureError( + "transfer backend returned a non-exception slot error" + ) + observed_error = error + elif state is WorkState.CANCELLED: + observed_error = TransferCancelledError("transfer slot was cancelled") + + # Error construction/provider inspection happens first; terminal state + # is the commit marker. A BaseException during error discovery leaves + # the prior nonterminal state, so the next observation retries instead + # of permanently caching a terminal state with no diagnostic. + if observed_state in _TERMINAL_STATES: + self._error = observed_error + self._state = observed_state + return observed_state + + def _raise_for_state(self) -> TransferSlot: + if self._state is WorkState.COMPLETED: + return self + if self._error is None: + self._error = BackendFailureError( + f"transfer slot reached {self._state.value} without an error" + ) + if isinstance(self._error, TransferError): + raise self._error + raise BackendFailureError("transfer slot failed") from self._error + + def _require_backend_slot(self) -> BackendRequestSlot: + backend_slot = self._backend_slot + if backend_slot is None: + raise BackendFailureError("transfer slot has no adopted backend request") + return backend_slot + + def _check_access(self) -> BoundTransfer: + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError("transfer slot is retained for cleanup only") + bound = self._bound + if bound is None: + raise ClosedError("transfer slot is closed") + owner_thread = self._owner_thread + if owner_thread is not None and threading.get_ident() != owner_thread: + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + # An open slot prevents its BoundTransfer, plan, and endpoint from + # closing. Creation validates that graph once, so polling need not walk + # it again on every request-state query. + return bound + + def __enter__(self) -> Self: + self._check_access() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class RequestLocalNotificationSlot(TransferSlot): + """Persistent request with an explicit notification per generation. + + The immutable operation, descriptors, index selection, CUDA-ordering + contract, and native request storage are inherited from ``TransferSlot``. + Only the attached notification varies. Every start must state either a + non-empty bytes-like payload or ``None``; the latter is an explicit native + clear, not "reuse the previous value". The distinct type and methods keep + all validation and dispatch branches out of fixed ``TransferSlot.execute``. + """ + + __slots__ = ( + "_backend_start_with_notification", + "_backend_start_and_poll_with_notification", + "_notification_fast_execute", + ) + + def __init__( + self, + endpoint: Endpoint, + bound: BoundTransfer, + *, + _token: object, + ) -> None: + super().__init__(endpoint, bound, _token=_token) + self._backend_start_with_notification: Callable[[bytes | None], None] | None = ( + None + ) + self._backend_start_and_poll_with_notification: ( + Callable[..., WorkState | None] | None + ) = None + self._notification_fast_execute: Callable[..., WorkState | None] | None = None + + def start( + self, + *, + notification: bytes | bytearray | memoryview | None, + ) -> None: + """Start one generation with an explicit notification or clear.""" + + if notification is not None: + if type(notification) is bytes: + if not notification: + raise ValueError( + "an attached notification payload must be non-empty" + ) + if len(notification) > _MAX_NOTIFICATION_BYTES: + raise ValueError( + "notification exceeds the " + f"{_MAX_NOTIFICATION_BYTES}-byte size limit" + ) + else: + notification = _validate_request_slot_notification(notification) + lock = self._lock + if lock is None: + self._start_with_notification_unlocked(notification) + return + with lock: + self._start_with_notification_unlocked(notification) + + def start_and_wait( + self, + *, + notification: bytes | bytearray | memoryview | None, + timeout: float | None = None, + ) -> RequestLocalNotificationSlot: + """Start with an explicit notification and wait for terminal state.""" + + if notification is not None: + if type(notification) is bytes: + if not notification: + raise ValueError( + "an attached notification payload must be non-empty" + ) + if len(notification) > _MAX_NOTIFICATION_BYTES: + raise ValueError( + "notification exceeds the " + f"{_MAX_NOTIFICATION_BYTES}-byte size limit" + ) + else: + notification = _validate_request_slot_notification(notification) + deadline_ns = _transfer_slot_deadline_ns(timeout) + start_and_poll = self._backend_start_and_poll_with_notification + poll_bounded = self._backend_poll_bounded + if start_and_poll is None: + lock = self._lock + if lock is None: + self._start_with_notification_unlocked(notification) + expected_generation = self._generation + else: + with lock: + self._start_with_notification_unlocked(notification) + expected_generation = self._generation + if poll_bounded is not None: + self._wait_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + else: + self._wait_until( + deadline_ns, + spin=True, + expected_generation=expected_generation, + ) + return self + + lock = self._lock + if lock is None: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, notification, deadline_ns + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self + if state in _TERMINAL_STATES: + return cast(RequestLocalNotificationSlot, self._raise_for_state()) + else: + with lock: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, notification, deadline_ns + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self + if state in _TERMINAL_STATES: + return cast(RequestLocalNotificationSlot, self._raise_for_state()) + + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if poll_bounded is not None: + self._wait_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + else: + self._wait_until( + deadline_ns, + spin=True, + expected_generation=expected_generation, + ) + return self + + def start_and_test( + self, + *, + notification: bytes | bytearray | memoryview | None, + ) -> bool: + """Start with an explicit tag/clear and consume immediate success.""" + + if notification is not None: + if type(notification) is bytes: + if not notification: + raise ValueError( + "an attached notification payload must be non-empty" + ) + if len(notification) > _MAX_NOTIFICATION_BYTES: + raise ValueError( + "notification exceeds the " + f"{_MAX_NOTIFICATION_BYTES}-byte size limit" + ) + else: + notification = _validate_request_slot_notification(notification) + + fast_start_and_test = self._notification_fast_execute + if fast_start_and_test is not None: + next_generation = self._generation + 1 + try: + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError( + "transfer slot is retained for cleanup only" + ) + if ( + self._owner_thread is not None + and _GET_IDENT() != self._owner_thread + ): + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + if self._active: + raise BusyError( + "transfer slot already owns an unrecycled generation" + ) + + if not self._caller_ready: + ready_event_synchronize = self._ready_event_synchronize + if ready_event_synchronize is not None: + try: + ready_event_synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to wait for CUDA readiness" + ) from error + else: + bound = self._bound + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + try: + self._generation = next_generation + self._active = True + except BaseException: + self._active = False + self._generation = next_generation - 1 + self._state = None + self._error = None + self._start_error = None + self._backend_generation_active = False + raise + + state = fast_start_and_test( + notification, + max_polls=0, + timeout_ns=None, + ) + except BaseException as error: + if self._generation != next_generation: + raise + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + if state is _WORK_COMPLETED: + self._active = False + return True + if state is _WORK_RUNNING: + # Match the fixed-notification slot's exact QD-many path: + # the provider already authenticated the enum and Core + # needs only to publish the acknowledged RUNNING state. + self._backend_generation_active = True + self._state = state + return False + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot notification start returned " + "without acknowledging a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + elif not isinstance(state, WorkState): + self._start_error = BackendFailureError( + "backend returned invalid fused request-slot state " + f"{state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + state = self._observe_backend_state_unlocked(state) + + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + start_and_poll = self._backend_start_and_poll_with_notification + lock = self._lock + if start_and_poll is None: + if lock is None: + return self._start_and_test_with_notification_fallback_unlocked( + notification + ) + with lock: + return self._start_and_test_with_notification_fallback_unlocked( + notification + ) + if lock is None: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, + notification, + None, + max_polls=0, + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + with lock: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, + notification, + None, + max_polls=0, + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def _start_and_test_with_notification_fallback_unlocked( + self, notification: bytes | None + ) -> bool: + """Keep tagged fallback start/observe/finalize under one slot lock.""" + + self._start_with_notification_unlocked(notification) + expected_generation = self._generation + state = self._refresh_unlocked() + if state is WorkState.COMPLETED: + self._finish_execute_success_unlocked(expected_generation) + return True + if state in _TERMINAL_STATES: + self._raise_for_state() + return False + + def execute( + self, + *, + notification: bytes | bytearray | memoryview | None, + timeout: float | None = None, + ) -> int: + """Execute and logically recycle one explicitly tagged generation.""" + + # Keep the common bytes/None cases inline: unlike the fixed slot this + # mode necessarily examines a request-local value, but it should not + # add a helper frame before the one fused pybind call. + if notification is not None: + if type(notification) is bytes: + if not notification: + raise ValueError( + "an attached notification payload must be non-empty" + ) + if len(notification) > _MAX_NOTIFICATION_BYTES: + raise ValueError( + "notification exceeds the " + f"{_MAX_NOTIFICATION_BYTES}-byte size limit" + ) + else: + notification = _validate_request_slot_notification(notification) + + fast_execute = self._notification_fast_execute + if timeout is None and fast_execute is not None: + next_generation = self._generation + 1 + try: + if self._closed: + raise ClosedError("transfer slot is closed") + if self._creation_failed: + raise BackendFailureError( + "transfer slot is retained for cleanup only" + ) + if ( + self._owner_thread is not None + and _GET_IDENT() != self._owner_thread + ): + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + if self._active: + raise BusyError( + "transfer slot already owns an unrecycled generation" + ) + + if not self._caller_ready: + ready_event_synchronize = self._ready_event_synchronize + if ready_event_synchronize is not None: + try: + ready_event_synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to wait for CUDA readiness" + ) from error + else: + bound = self._bound + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + try: + self._generation = next_generation + self._active = True + except BaseException: + self._active = False + self._generation = next_generation - 1 + self._state = None + self._error = None + self._start_error = None + self._backend_generation_active = False + raise + + state = fast_execute( + notification, + max_polls=_FUSED_REQUEST_SLOT_MAX_POLLS, + timeout_ns=None, + ) + except BaseException as error: + if self._generation != next_generation: + raise + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + if state is _WORK_COMPLETED: + self._active = False + return next_generation + if state is _WORK_RUNNING: + # Preserve the same ownership-before-state publication as + # the fixed slot. Disable the request-local dispatcher + # before the inherited continuation: an interruption after + # its idle commit must not leave a stale fast dispatcher. + self._backend_generation_active = True + self._state = state + self._notification_fast_execute = None + poll_bounded = self._backend_poll_bounded + if poll_bounded is not None: + return self._wait_execute_bounded_until( + None, + expected_generation=next_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + None, expected_generation=next_generation + ) + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot notification start returned " + "without acknowledging a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + elif not isinstance(state, WorkState): + self._start_error = BackendFailureError( + "backend returned invalid fused request-slot state " + f"{state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + state = refreshed + else: + state = self._observe_backend_state_unlocked(state) + + expected_generation = self._generation + if state in _TERMINAL_STATES: + self._raise_for_state() + # A bounded continuation finalizes through the inherited idle + # transaction. Disable this subclass's exact dispatcher first so + # an interruption after that transaction's public idle commit + # cannot pair it with a stale terminal state. + self._notification_fast_execute = None + poll_bounded = self._backend_poll_bounded + if poll_bounded is not None: + return self._wait_execute_bounded_until( + None, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + None, + expected_generation=expected_generation, + ) + + deadline_ns = None if timeout is None else _transfer_slot_deadline_ns(timeout) + start_and_poll = self._backend_start_and_poll_with_notification + lock = self._lock + poll_bounded = self._backend_poll_bounded + if start_and_poll is None: + if lock is None: + self._start_with_notification_unlocked(notification) + expected_generation = self._generation + else: + with lock: + self._start_with_notification_unlocked(notification) + expected_generation = self._generation + if poll_bounded is not None: + return self._wait_execute_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + deadline_ns, + expected_generation=expected_generation, + ) + + if lock is None: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, + notification, + deadline_ns, + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self._finish_execute_success_unlocked(expected_generation) + if state in _TERMINAL_STATES: + self._raise_for_state() + else: + with lock: + state = self._start_and_poll_with_notification_unlocked( + start_and_poll, + notification, + deadline_ns, + ) + expected_generation = self._generation + if state is WorkState.COMPLETED: + return self._finish_execute_success_unlocked(expected_generation) + if state in _TERMINAL_STATES: + self._raise_for_state() + + if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: + self._validate_wait_generation(expected_generation) + raise TransferTimeoutError("transfer slot timed out") + if poll_bounded is not None: + return self._wait_execute_bounded_until( + deadline_ns, + expected_generation=expected_generation, + poll_bounded=poll_bounded, + ) + return self._wait_execute_until( + deadline_ns, + expected_generation=expected_generation, + ) + + def _start_with_notification_unlocked(self, notification: bytes | None) -> None: + previous_generation = self._generation + try: + self._begin_generation_unlocked() + self._notification_fast_execute = None + start = self._backend_start_with_notification + if start is None: + self._start_error = BackendFailureError( + "transfer slot has no notification-override operation" + ) + return + if self._lock is None: + result = start(notification) + else: + result = self._endpoint._backend_request_slot_start_with_notification( + start, notification + ) + self._backend_generation_active = True + if result is not None: + raise BackendFailureError( + "backend request-slot start_with_notification must return None" + ) + except BaseException as error: + if self._generation == previous_generation: + raise + self._start_error = error + + def _start_and_poll_with_notification_unlocked( + self, + start_and_poll: Callable[..., WorkState | None], + notification: bytes | None, + deadline_ns: int | None, + *, + max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, + ) -> WorkState: + previous_generation = self._generation + try: + self._begin_generation_unlocked() + self._notification_fast_execute = None + timeout_ns = ( + None + if deadline_ns is None + else min( + _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS, + max(0, deadline_ns - time.monotonic_ns()), + ) + ) + if self._lock is None: + state = start_and_poll( + notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + else: + state = self._endpoint._backend_request_slot_start_and_poll_with_notification( + start_and_poll, + notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + except BaseException as error: + if self._generation == previous_generation: + raise + self._start_error = error + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + if state is None: + self._start_error = BackendFailureError( + "backend fused request-slot notification start returned without " + "acknowledging a generation" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + if not isinstance(state, WorkState): + self._start_error = BackendFailureError( + f"backend returned invalid fused request-slot state {state!r}" + ) + refreshed = self._refresh_unlocked() + assert refreshed is not None + return refreshed + return self._observe_backend_state_unlocked(state) + + def _recycle_unlocked(self) -> None: + """Disable notification fast dispatch before an exceptional recycle.""" + + # Do not penalize or mutate an erroneous idle recycle. Active failure + # and cancellation recovery, however, may publish public idle before + # clearing retained diagnostics in the base transaction. Keeping the + # notification dispatcher disabled makes that dirty-idle boundary a + # safe generic fallback. + if self._active: + self._notification_fast_execute = None + super()._recycle_unlocked() + + def _restore_fast_execute_unlocked(self) -> None: + """Restore only the notification-aware exact-execute callable.""" + + self._fast_execute = None + start_and_poll = self._backend_start_and_poll_with_notification + self._notification_fast_execute = ( + start_and_poll + if ( + self._lock is None + and start_and_poll is not None + and self._backend_deferred_successful_recycle + ) + else None + ) + + def _adopt_backend_slot_unlocked( + self, + backend_slot: BackendRequestSlot, + *, + execution_dispatch: object | None, + ) -> None: + dispatch = backend_slot if execution_dispatch is None else execution_dispatch + start = getattr(dispatch, "start_with_notification", None) + start_and_poll = ( + getattr(dispatch, "start_and_poll_with_notification", None) + if self._endpoint.capabilities.fused_request_slot_start_poll + else None + ) + if not callable(start): + raise BackendFailureError( + "backend request-slot start_with_notification operation must be " + "callable" + ) + if ( + getattr(start, "__func__", None) + is BackendRequestSlot.start_with_notification + ): + raise BackendFailureError( + "backend advertised request-slot notification overrides without " + "implementing start_with_notification" + ) + if self._endpoint.capabilities.fused_request_slot_start_poll and not callable( + start_and_poll + ): + raise BackendFailureError( + "backend request-slot start_and_poll_with_notification operation " + "must be callable" + ) + if ( + start_and_poll is not None + and getattr(start_and_poll, "__func__", None) + is BackendRequestSlot.start_and_poll_with_notification + ): + raise BackendFailureError( + "backend advertised fused request-slot notification overrides " + "without implementing start_and_poll_with_notification" + ) + + # Publish the extra dispatch table before the base adoption commit. + # An interruption can therefore expose either no backend slot or one + # with every method required by its public dynamic type. + self._backend_start_with_notification = start + self._backend_start_and_poll_with_notification = start_and_poll + super()._adopt_backend_slot_unlocked( + backend_slot, + execution_dispatch=execution_dispatch, + ) + # Even an explicit unbound TransferSlot.start(dynamic_slot) must not + # enter the fixed provider family: NIXL's legacy nullptr options mean + # "retain the last tag", not "clear". Leave only observation, + # recycle, cancellation, and close shared between the two families. + self._backend_start = None + self._backend_start_and_poll = None + self._fast_execute = None + self._restore_fast_execute_unlocked() + + +class Work: + """Request-local asynchronous transfer state.""" + + __slots__ = ( + "_backend_work", + "_backend_released", + "_bound", + "_closed", + "_dependencies", + "_endpoint", + "_error", + "_owned_plan", + "_plan", + "_recovery_token", + "_lock", + "_owner_thread", + "_submission_error", + "_submitting", + "_state", + ) + + def __init__( + self, + endpoint: Endpoint, + backend_work: BackendWork | None, + *, + plan: TransferPlan | None = None, + bound: BoundTransfer | None = None, + dependencies: tuple[object, ...] = (), + owned_plan: bool = False, + recovery_token: object | None = None, + _token: object, + ) -> None: + if _token is not _CONSTRUCTION_TOKEN: + raise TypeError("Work objects are created by Endpoint operations") + self._endpoint = endpoint + self._backend_work = backend_work + self._backend_released = False + self._plan = plan + self._bound = bound + self._dependencies = dependencies + self._owned_plan = owned_plan + # Passive caller identity: never interpret or provider-forward it. + self._recovery_token = recovery_token + self._lock = ( + None + if endpoint._thread_mode + in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) + else threading.RLock() + ) + self._owner_thread = ( + endpoint._owner_thread + if endpoint._thread_mode is ThreadMode.SINGLE + else None + ) + self._submission_error: BaseException | None = None + # A Work is published before provider entry so an interrupted caller can + # recover its identity. Observation is unsafe until submission return or + # exception reconciliation has completed. + self._submitting = True + self._state = WorkState.PENDING + self._error: BaseException | None = None + self._closed = False + + @property + def state(self) -> WorkState: + return self._refresh() + + @property + def recovery_token(self) -> object | None: + """Return the opaque identity retained for lost-return recovery. + + Match this object with ``is`` in :attr:`Endpoint.open_works`. Core + never hashes, compares, invokes, serializes, or passes it to a provider. + """ + + return self._recovery_token + + @property + def error(self) -> BaseException | None: + self._refresh() + return self._error + + def exception(self) -> BaseException | None: + return self.error + + @property + def completion_event(self) -> object | None: + lock = self._lock + if lock is None: + return self._completion_event_unlocked() + with lock: + return self._completion_event_unlocked() + + def _completion_event_unlocked(self) -> object | None: + self._check_not_closed() + self._check_submission_complete_unlocked() + if not self._endpoint.capabilities.cuda_completion_event: + return None + if self._backend_work is None: + raise BusyError("transfer work has not been adopted by its backend") + return self._endpoint._backend_call( + "query work completion event", + lambda: self._backend_work.completion_event, + ) + + def _adopt_backend_work(self, backend_work: BackendWork) -> None: + """Adopt provider state without allocating a submission closure.""" + + if not isinstance(backend_work, BackendWork): + raise BackendFailureError("backend must adopt a BackendWork instance") + lock = self._lock + if lock is None: + self._adopt_backend_work_unlocked(backend_work) + return + with lock: + self._adopt_backend_work_unlocked(backend_work) + + def _adopt_backend_work_unlocked(self, backend_work: BackendWork) -> None: + if self._closed: + raise BackendFailureError( + "backend attempted to adopt an expired transfer work" + ) + if self._backend_work is not None: + raise BackendFailureError( + "backend attempted to adopt transfer work more than once" + ) + self._backend_work = backend_work + + @property + def closed(self) -> bool: + return self._closed + + def test(self) -> bool: + return self._refresh() in _TERMINAL_STATES + + def poll_bounded( + self, + *, + max_polls: int = _FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + timeout_ns: int | None = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, + timeout_check_interval: int = (_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL), + ) -> WorkState: + """Perform a bounded poll and return the exact observed state. + + A provider advertising ``fused_work_poll`` performs the entire loop in + one provider call. Other providers retain identical semantics through + repeated ordinary state observations. The returned state may be + nonterminal; this Work remains the authoritative owner. + """ + + _validate_bounded_poll_arguments( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + allow_zero=False, + ) + if self._endpoint.capabilities.fused_work_poll: + lock = self._lock + if lock is None: + return self._poll_bounded_unlocked( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + with lock: + return self._poll_bounded_unlocked( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + + deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns + state = self._state + for poll_index in range(max_polls): + state = self._refresh() + if state in _TERMINAL_STATES: + return state + if ( + deadline_ns is not None + and (poll_index == 0 or (poll_index + 1) % timeout_check_interval == 0) + and time.monotonic_ns() >= deadline_ns + ): + break + return state + + def _poll_bounded_unlocked( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState: + self._check_not_closed() + self._check_thread_unlocked() + self._check_submission_complete_unlocked() + if self._state in _TERMINAL_STATES: + return self._state + backend_work = self._backend_work + if backend_work is None: + return self._state + state = self._endpoint._backend_work_poll_bounded( + backend_work, + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + self._ingest_backend_state_unlocked(state) + return self._state + + def test_bounded( + self, + *, + max_polls: int = _FUSED_REQUEST_SLOT_POLL_MAX_POLLS, + timeout_ns: int | None = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, + timeout_check_interval: int = (_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL), + ) -> bool: + """Compatibility wrapper returning whether a bounded poll is terminal.""" + + return ( + self.poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + in _TERMINAL_STATES + ) + + def wait(self, timeout: float | None = None) -> Work: + if timeout is not None: + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): + raise TypeError("timeout must be a non-negative number or None") + if not math.isfinite(timeout) or timeout < 0: + raise ValueError("timeout must be finite and non-negative") + deadline = None if timeout is None else time.monotonic() + timeout + polls = 0 + while True: + state = self._refresh() + if state in _TERMINAL_STATES: + return self._raise_for_state() + if deadline is not None and time.monotonic() >= deadline: + raise TransferTimeoutError("transfer work timed out") + if self._endpoint.progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + polls += 1 + if polls >= 16: + # Yield without imposing a fixed millisecond latency floor. + # Long waits back off to cap host polling overhead. + time.sleep(0 if polls < 128 else 0.0001) + + def result(self, timeout: float | None = None) -> None: + self.wait(timeout) + return None + + def cancel(self) -> bool: + lock = self._lock + if lock is None: + return self._cancel_unlocked() + with lock: + return self._cancel_unlocked() + + def _cancel_unlocked(self) -> bool: + self._check_not_closed() + self._check_submission_complete_unlocked() + if self._refresh_unlocked() in _TERMINAL_STATES: + return self._state is WorkState.CANCELLED + if self._backend_work is None: + return False + cancelled = self._endpoint._backend_call( + "cancel work", self._backend_work.cancel + ) + if not isinstance(cancelled, bool): + raise BackendFailureError("backend work cancellation must return a bool") + self._refresh_unlocked() + return cancelled and self._state is WorkState.CANCELLED + + def close(self) -> None: + state_lock = self._endpoint._state_lock + if state_lock is _NOOP_LOCK: + work_lock = self._lock + if work_lock is None: + self._close_unlocked() + return + with work_lock: + self._close_unlocked() + return + with state_lock: + self._close_with_work_lock() + + def _close_with_work_lock(self) -> None: + lock = self._lock + if lock is None: + self._close_unlocked() + return + with lock: + self._close_unlocked() + + def _close_unlocked(self) -> None: + if self._closed: + # CLOSED is the provider-release commit, not proof that every + # Core-only graph edge and retained reference was cleared. A + # retry must finish an epilogue interrupted after that commit. + self._finish_close_unlocked() + return + self._check_submission_complete_unlocked() + self._check_thread_unlocked() + if ( + self._state not in _TERMINAL_STATES + and self._refresh_unlocked() not in _TERMINAL_STATES + ): + raise BusyError("cannot close non-terminal transfer work") + if not self._backend_released: + if self._backend_work is None: + raise BackendFailureError( + "terminal transfer work has no backend request" + ) + self._endpoint._backend_work_close(self._backend_work) + self._backend_released = True + if self._plan is not None: + self._plan._works.discard(self) + if self._bound is not None: + self._bound._works.discard(self) + try: + if self._owned_plan and self._plan is not None: + self._plan.close() + except BaseException: + if self._plan is not None: + self._plan._works.add(self) + raise + self._closed = True + self._finish_close_unlocked() + + def _finish_close_unlocked(self) -> None: + """Idempotently unlink a closed Work from the Core lifetime graph.""" + + plan = self._plan + bound = self._bound + # Retain both parent references until every blocking edge is gone. An + # interrupt during any discard therefore leaves enough state to retry. + if plan is not None: + plan._works.discard(self) + if bound is not None: + bound._works.discard(self) + self._endpoint._works.discard(self) + self._dependencies = () + self._plan = None + self._bound = None + self._backend_work = None + self._error = None + self._submission_error = None + # Commit last: absence means every provider/Core lifetime edge is gone. + self._recovery_token = None + + async def _wait_async(self) -> Work: + while True: + state = self._refresh() + if state in _TERMINAL_STATES: + return self._raise_for_state() + if self._endpoint.progress_mode is ProgressMode.MANUAL: + self._endpoint.progress() + await asyncio.sleep(0) + + def __await__(self): # type: ignore[no-untyped-def] + return self._wait_async().__await__() + + def _refresh(self) -> WorkState: + lock = self._lock + if lock is None: + return self._refresh_unlocked() + with lock: + return self._refresh_unlocked() + + def _refresh_unlocked(self) -> WorkState: + self._check_not_closed() + self._check_thread_unlocked() + self._check_submission_complete_unlocked() + if self._state in _TERMINAL_STATES: + return self._state + if self._backend_work is None: + return self._state + state = self._endpoint._backend_work_state(self._backend_work) + return self._ingest_backend_state_unlocked(state) + + def _ingest_backend_state_unlocked(self, state: WorkState) -> WorkState: + """Validate and cache a provider state already observed this call.""" + + if not isinstance(state, WorkState): + raise BackendFailureError(f"backend returned invalid work state {state!r}") + if self._state is WorkState.RUNNING and state is WorkState.PENDING: + raise BackendFailureError( + "backend work state regressed from running to pending" + ) + observed_state = state + observed_error = self._error + if state in _TERMINAL_STATES and self._submission_error is not None: + observed_state = WorkState.FAILED + observed_error = self._submission_error + elif state is WorkState.FAILED: + backend_work = self._backend_work + if backend_work is None: + observed_error = BackendFailureError( + "failed transfer work has no backend request" + ) + else: + try: + error = self._endpoint._backend_call( + "query work error", lambda: backend_work.error + ) + except Exception as observation_error: + observed_error = observation_error + else: + if error is None: + error = BackendFailureError( + "transfer backend reported failure without an error" + ) + elif not isinstance(error, BaseException): + error = BackendFailureError( + "transfer backend returned a non-exception error" + ) + observed_error = error + elif state is WorkState.CANCELLED: + observed_error = TransferCancelledError("transfer work was cancelled") + + # Terminal state is the public commit. Derive its diagnostic first so + # an asynchronous BaseException during provider error observation leaves + # the prior nonterminal state retryable. + if observed_state in _TERMINAL_STATES: + self._error = observed_error + self._state = observed_state + return self._state + + def _raise_for_state(self) -> Work: + if self._state is WorkState.COMPLETED: + return self + if self._error is None: + self._error = BackendFailureError( + f"transfer work reached {self._state.value} without an error" + ) + if isinstance(self._error, TransferError): + raise self._error + raise BackendFailureError("transfer work failed") from self._error + + def _check_not_closed(self) -> None: + if self._closed: + raise ClosedError("transfer work is closed") + + def _check_submission_complete_unlocked(self) -> None: + if self._submitting: + raise BusyError("transfer work submission is in progress") + + def _check_thread_unlocked(self) -> None: + owner_thread = self._owner_thread + if owner_thread is not None and _GET_IDENT() != owner_thread: + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + + +class Endpoint: + """Backend-neutral endpoint for asymmetric one-sided data transfer.""" + + def __init__( + self, + name: str, + *, + backend: str = "reference", + endpoint_id: str | None = None, + progress_mode: ProgressMode | str = ProgressMode.MANUAL, + thread_mode: ThreadMode | str = ThreadMode.SERIALIZED, + options: Mapping[str, Any] | None = None, + ) -> None: + _validate_nonempty_str(name, "name") + _validate_nonempty_str(backend, "backend") + if endpoint_id is not None: + _validate_nonempty_str(endpoint_id, "endpoint_id") + progress_mode = _coerce_enum(progress_mode, ProgressMode, "progress_mode") + thread_mode = _coerce_enum(thread_mode, ThreadMode, "thread_mode") + if options is None: + options = {} + elif not isinstance(options, Mapping): + raise TypeError("options must be a mapping or None") + + self._name = name + self._backend_name = backend + self._endpoint_id = endpoint_id or str(uuid.uuid4()) + self._incarnation = str(uuid.uuid4()) + self._progress_mode = progress_mode + self._thread_mode = thread_mode + self._owner_thread = threading.get_ident() + self._backend_lock = ( + threading.RLock() if thread_mode is ThreadMode.SERIALIZED else _NOOP_LOCK + ) + self._state_lock = ( + _NOOP_LOCK + if thread_mode in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) + else threading.RLock() + ) + self._notification_consumer_lock = ( + threading.RLock() if thread_mode is ThreadMode.MULTIPLE else _NOOP_LOCK + ) + self._active_multiple_backend_calls: ( + set[weakref.ReferenceType[_BackendCallReservation]] | None + ) = (set() if thread_mode is ThreadMode.MULTIPLE else None) + self._multiple_backend_call_reservation_pool: ( + list[_BackendCallReservation] | None + ) = ([] if thread_mode is ThreadMode.MULTIPLE else None) + self._registrations: dict[str, Registration] = {} + self._registration_names: dict[str, Registration] = {} + self._peers: dict[str, Peer] = {} + self._plans: set[TransferPlan] = set() + self._orphaned_plans: set[TransferPlan] = set() + self._bound_transfers: set[BoundTransfer] = set() + self._transfer_slots: set[TransferSlot] = set() + self._works: set[Work] = set() + self._closing = False + self._creation_failed = False + self._closed = False + + _retry_orphaned_backend_receipts(backend, self._endpoint_id) + registration = _get_backend_registration(backend) + backend_options = MappingProxyType(dict(options)) + receipt = _BackendConstructionReceipt(backend, self._endpoint_id) + self._backend_receipt = receipt + self._backend_finalizer = weakref.finalize( + self, _finalize_backend_receipt, receipt + ) + self._backend_finalizer.atexit = False + endpoint_ref = weakref.ref(self) + + def adopt_backend(implementation: Backend) -> None: + # The receipt owns the backend before an Endpoint attribute becomes + # observable. The outer transaction repairs either STORE_ATTR cut. + receipt.adopt(implementation) + endpoint = endpoint_ref() + if endpoint is None: + raise BackendFailureError( + "endpoint disappeared while adopting its backend" + ) + endpoint._backend = implementation + endpoint._creation_failed = True + + factory = registration.factory + unowned_v1_result = False + try: + if registration.factory_api_version == 1: + candidate = cast(Callable[..., object], factory)( + self._endpoint_id, + name, + self._incarnation, + progress_mode, + thread_mode, + backend_options, + ) + if not isinstance(candidate, Backend): + # A wrong-type legacy return was never adopted by Core. It + # may be an arbitrary foreign object, so do not call close + # or inspect any provider-shaped attribute on it. + unowned_v1_result = True + raise TypeError("backend factory must return a Backend instance") + # This compatibility path necessarily retains the legacy + # CALL-to-adoption gap. Version-2 providers close it. + adopt_backend(candidate) + implementation = receipt.seal_successful_call(None) + else: + result = cast(Callable[..., object], factory)( + self._endpoint_id, + name, + self._incarnation, + progress_mode, + thread_mode, + backend_options, + adopt_backend=adopt_backend, + ) + implementation = receipt.seal_successful_call(result) + except BaseException as error: + receipt.seal_failed_call() + try: + self._cleanup_adopted_backend_after_creation_error(error) + except BaseException: + raise + if unowned_v1_result: + raise + if isinstance(error, TransferError) or not isinstance(error, Exception): + raise + raise BackendFailureError( + f"failed to create transfer backend {backend!r}" + ) from error + + try: + # From this assignment onward Core owns the returned backend. Keep + # capability discovery, dispatch caching, and EndpointInfo + # publication in one cleanup transaction. + self._backend = implementation + self._creation_failed = True + capabilities = implementation.capabilities + if not isinstance(capabilities, Capabilities): + raise TypeError("backend capabilities must be Capabilities") + if ( + progress_mode is ProgressMode.MANUAL + and not capabilities.manual_progress + ): + raise UnsupportedError( + f"backend {backend!r} does not support manual progress" + ) + if ( + progress_mode is ProgressMode.BACKGROUND + and not capabilities.background_progress + ): + raise UnsupportedError( + f"backend {backend!r} does not support background progress" + ) + supported_thread_modes = implementation.supported_thread_modes + if not isinstance(supported_thread_modes, frozenset) or not all( + isinstance(mode, ThreadMode) for mode in supported_thread_modes + ): + raise TypeError( + "backend supported_thread_modes must be a frozenset of " + "ThreadMode values" + ) + if thread_mode not in supported_thread_modes: + raise UnsupportedError( + f"backend {backend!r} does not support thread mode " + f"{thread_mode.value!r}" + ) + send_notification_sync_dispatch = ( + cast( + Callable[[object, bytes], object | None], + implementation.send_notification_sync, + ) + if capabilities.synchronous_notification_send + else None + ) + info = EndpointInfo( + backend=backend, + endpoint_id=self._endpoint_id, + name=name, + incarnation=self._incarnation, + capabilities=capabilities, + ) + self._capabilities = capabilities + self._send_notification_sync_dispatch = send_notification_sync_dispatch + self._info = info + self._creation_failed = False + receipt.mark_published() + except BaseException as error: + # Re-establish both fields if interruption followed either + # STORE_ATTR. Failed cleanup is exposed through a recovery endpoint. + self._backend = implementation + self._creation_failed = True + try: + self._close_failed_endpoint_backend() + except BaseException as cleanup_error: + raise _endpoint_creation_cleanup_failure( + self, error, cleanup_error + ) from cleanup_error + raise + + @property + def name(self) -> str: + return self._name + + @property + def backend(self) -> str: + return self._backend_name + + @property + def endpoint_id(self) -> str: + return self._endpoint_id + + @property + def incarnation(self) -> str: + return self._incarnation + + @property + def identity(self) -> EndpointIdentity: + return EndpointIdentity(self.endpoint_id, self.incarnation) + + @property + def progress_mode(self) -> ProgressMode: + return self._progress_mode + + @property + def thread_mode(self) -> ThreadMode: + return self._thread_mode + + @property + def capabilities(self) -> Capabilities: + return self._capabilities + + @property + def info(self) -> EndpointInfo: + return self._info + + @property + def closed(self) -> bool: + return self._closed + + @property + def open_registrations(self) -> tuple[Registration, ...]: + """Return registrations owned by this endpoint for recovery/inspection.""" + + with self._state_lock: + self._check_open() + return tuple(self._registrations.values()) + + @property + def open_peers(self) -> tuple[Peer, ...]: + """Return imported peers owned by this endpoint for recovery/inspection.""" + + with self._state_lock: + self._check_open() + return tuple(self._peers.values()) + + @property + def open_plans(self) -> tuple[TransferPlan, ...]: + """Return plans owned by this endpoint for recovery/inspection.""" + + with self._state_lock: + self._check_open() + return tuple(self._plans) + + @property + def open_works(self) -> tuple[Work, ...]: + """Return open Work objects, including interrupted submissions. + + The snapshot lets an application recover and drain a Work if an + asynchronous exception lands after submission but before the caller + stores the returned object. + """ + + with self._state_lock: + self._check_open() + return tuple(self._works) + + @property + def open_bound_transfers(self) -> tuple[BoundTransfer, ...]: + """Return reusable bindings owned by this endpoint for recovery.""" + + with self._state_lock: + self._check_open() + return tuple(self._bound_transfers) + + @property + def open_transfer_slots(self) -> tuple[TransferSlot, ...]: + """Return persistent request slots owned by this endpoint for recovery.""" + + with self._state_lock: + self._check_open() + return tuple(self._transfer_slots) + + @_state_locked + def register( + self, + tensor_or_storage: object, + *, + name: str | None = None, + ) -> Registration: + self._check_open() + if name is None: + name = f"registration-{uuid.uuid4()}" + _validate_nonempty_str(name, "name") + if name in self._registration_names: + raise InvalidRegionError(f"registration name {name!r} is not unique") + registration_id = str(uuid.uuid4()) + memory = _registered_memory(tensor_or_storage, registration_id, name) + for existing in self._registrations.values(): + _validate_memory_stable(existing._memory) + if _memory_ranges_overlap(memory, existing._memory): + raise InvalidRegionError( + f"registration {name!r} overlaps existing registration " + f"{existing.name!r}" + ) + if isinstance(tensor_or_storage, RawSpan) and not self.capabilities.raw_spans: + raise UnsupportedError( + f"backend {self.backend!r} does not support raw spans" + ) + registration = Registration( + self, + memory, + _UNADOPTED_HANDLE, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._registrations[registration_id] = registration + self._registration_names[name] = registration + result = self._backend_call( + "register memory", + lambda: self._backend.register( + memory, + adopt_handle=lambda handle: _adopt_resource_handle( + registration, handle, "registration" + ), + ), + ) + if result is not None: + raise BackendFailureError( + "backend register must return None after adopting its handle" + ) + if registration._handle is _UNADOPTED_HANDLE: + raise BackendFailureError( + "backend register returned without adopting a handle" + ) + return registration + except BaseException as error: + if registration._handle is not _UNADOPTED_HANDLE: + try: + registration.close() + except BaseException as cleanup_error: + registration._creation_failed = True + raise _creation_cleanup_failure( + "registration", error, cleanup_error + ) from cleanup_error + else: + self._registrations.pop(registration_id, None) + self._registration_names.pop(name, None) + registration._closed = True + registration._handle = None + registration._memory = _released_memory(memory) + raise + + @_state_locked + def export_metadata( + self, + registrations: Sequence[Registration] | None = None, + ) -> EndpointMetadata: + self._check_open() + selected = ( + tuple(self._registrations.values()) + if registrations is None + else tuple(registrations) + ) + if len(selected) != len(set(selected)): + raise InvalidMetadataError("registrations must not contain duplicates") + values: list[RegistrationMetadata] = [] + for registration in selected: + self._validate_registration(registration) + descriptor = self._backend_call( + "describe registration", + lambda registration=registration: ( + self._backend.describe_registration(registration._handle) + ), + ) + if not isinstance(descriptor, (bytes, bytearray, memoryview)): + raise BackendFailureError( + "backend registration descriptor must be bytes-like" + ) + values.append( + RegistrationMetadata( + registration_id=registration.registration_id, + name=registration.name, + nbytes=registration.nbytes, + device=registration.device, + memory_type=registration.memory_type, + backend_descriptor=bytes(descriptor), + ) + ) + payload = self._backend_call( + "export endpoint metadata", + lambda: self._backend.endpoint_payload( + tuple(registration._handle for registration in selected) + ), + ) + if not isinstance(payload, (bytes, bytearray, memoryview)): + raise BackendFailureError("backend endpoint payload must be bytes-like") + return EndpointMetadata( + wire_major=EndpointMetadata.WIRE_MAJOR, + wire_minor=EndpointMetadata.WIRE_MINOR, + required_features=(), + backend=self.backend, + endpoint_id=self.endpoint_id, + name=self.name, + incarnation=self.incarnation, + backend_payload=bytes(payload), + registrations=tuple(values), + ) + + @_state_locked + def import_peer( + self, + metadata: EndpointMetadata | bytes | bytearray | memoryview, + *, + expected_identity: EndpointIdentity | None = None, + expected_endpoint_id: str | None = None, + expected_name: str | None = None, + expected_incarnation: str | None = None, + ) -> Peer: + self._check_open() + if not isinstance(metadata, EndpointMetadata): + metadata = EndpointMetadata.from_bytes(metadata) + if expected_identity is not None: + if not isinstance(expected_identity, EndpointIdentity): + raise TypeError("expected_identity must be an EndpointIdentity") + if expected_endpoint_id is not None or expected_incarnation is not None: + raise ValueError( + "expected_identity is mutually exclusive with separate expected " + "endpoint_id/incarnation values" + ) + expected_endpoint_id = expected_identity.endpoint_id + expected_incarnation = expected_identity.incarnation + expected_fields = { + "endpoint_id": expected_endpoint_id, + "name": expected_name, + "incarnation": expected_incarnation, + } + for field_name, expected in expected_fields.items(): + if expected is None: + continue + _validate_nonempty_str(expected, f"expected_{field_name}") + actual = getattr(metadata, field_name) + if actual != expected: + raise InvalidMetadataError( + f"peer {field_name} {actual!r} does not match expected " + f"value {expected!r}" + ) + if metadata.backend != self.backend: + raise BackendMismatchError( + f"peer backend {metadata.backend!r} does not match " + f"local backend {self.backend!r}" + ) + existing = self._peers.get(metadata.endpoint_id) + if existing is not None: + if existing.incarnation != metadata.incarnation: + raise BusyError( + f"peer {metadata.endpoint_id!r} changed incarnation; " + "close the old peer before importing the replacement" + ) + raise BusyError(f"peer {metadata.endpoint_id!r} is already imported") + peer = Peer( + self, + metadata, + _UNADOPTED_HANDLE, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._peers[peer.endpoint_id] = peer + result = self._backend_call( + "import peer", + lambda: self._backend.import_peer( + metadata, + adopt_handle=lambda handle: _adopt_resource_handle( + peer, handle, "peer" + ), + ), + ) + if result is not None: + raise BackendFailureError( + "backend peer import must return None after adopting its handle" + ) + if peer._handle is _UNADOPTED_HANDLE: + raise BackendFailureError( + "backend peer import returned without adopting a handle" + ) + return peer + except BaseException as error: + if peer._handle is not _UNADOPTED_HANDLE: + try: + peer.close() + except BaseException as cleanup_error: + peer._creation_failed = True + raise _creation_cleanup_failure( + "peer", error, cleanup_error + ) from cleanup_error + else: + self._peers.pop(peer.endpoint_id, None) + peer._closed = True + peer._handle = None + peer._registrations = {} + raise + + @_state_locked + def prepare( + self, + *, + local: Sequence[Region], + remote: Sequence[Region], + indexed: bool = False, + ) -> TransferPlan: + self._check_open() + if not isinstance(indexed, bool): + raise TypeError("indexed must be a bool") + local = tuple(local) + remote = tuple(remote) + self._validate_regions(local, remote, indexed=indexed) + backend_local = tuple(region._to_backend() for region in local) + backend_remote = tuple(region._to_backend() for region in remote) + # Finish every Python allocation and publish the local lifetime guard + # before asking a provider to create native state. + plan = TransferPlan( + self, + _UNADOPTED_HANDLE, + local, + remote, + indexed, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._plans.add(plan) + result = self._backend_call( + "prepare transfer plan", + lambda: self._backend.prepare( + backend_local, + backend_remote, + indexed=indexed, + adopt_handle=lambda handle: _adopt_resource_handle( + plan, handle, "transfer plan" + ), + ), + ) + if result is not None: + raise BackendFailureError( + "backend prepare must return None after adopting its handle" + ) + if plan._handle is _UNADOPTED_HANDLE: + raise BackendFailureError( + "backend prepare returned without adopting a handle" + ) + return plan + except BaseException as error: + if plan._handle is not _UNADOPTED_HANDLE: + try: + plan.close() + except BaseException as cleanup_error: + plan._creation_failed = True + self._orphaned_plans.add(plan) + raise _creation_cleanup_failure( + "transfer plan", error, cleanup_error + ) from cleanup_error + else: + self._plans.discard(plan) + plan._closed = True + plan._handle = None + plan._local = () + plan._remote = () + plan._origins = frozenset() + plan._reuse_tokens = {} + raise + + def read( + self, + plan: TransferPlan | None = None, + *, + local: Sequence[Region] | None = None, + remote: Sequence[Region] | None = None, + local_indices: Sequence[int] | None = None, + remote_indices: Sequence[int] | None = None, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + ) -> Work: + return self._submit( + TransferOp.READ, + plan, + local=local, + remote=remote, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + + def write( + self, + plan: TransferPlan | None = None, + *, + local: Sequence[Region] | None = None, + remote: Sequence[Region] | None = None, + local_indices: Sequence[int] | None = None, + remote_indices: Sequence[int] | None = None, + notification: bytes | bytearray | memoryview | None = None, + cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, + stream: object | None = None, + ready_event: object | None = None, + recovery_token: object | None = None, + ) -> Work: + return self._submit( + TransferOp.WRITE, + plan, + local=local, + remote=remote, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + + @_state_locked + def _select_plan_indices( + self, + plan: TransferPlan, + op: TransferOp | str, + *, + local_indices: Sequence[int], + remote_indices: Sequence[int], + ) -> IndexSelection: + self._check_open() + self._validate_plan_locked(plan) + if not plan.indexed: + raise InvalidRegionError( + "indices require a plan prepared with indexed=True" + ) + op = _coerce_enum(op, TransferOp, "op") + supported = ( + self.capabilities.read if op is TransferOp.READ else self.capabilities.write + ) + if not supported: + raise UnsupportedError( + f"backend {self.backend!r} does not support {op.value}" + ) + backend_selection, local_devices = _snapshot_backend_index_selection( + plan, op, local_indices, remote_indices + ) + return IndexSelection( + self, + plan, + op, + backend_selection, + local_devices, + _token=_CONSTRUCTION_TOKEN, + ) + + @_state_locked + def _bind_transfer( + self, + plan: TransferPlan, + op: TransferOp | str, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | bytearray | memoryview | None, + cuda_ordering: CudaOrderingMode | str, + stream: object | None, + ready_event: object | None, + ) -> BoundTransfer: + self._check_open() + self._validate_plan_locked(plan) + op = _coerce_enum(op, TransferOp, "op") + supported = ( + self.capabilities.read if op is TransferOp.READ else self.capabilities.write + ) + if not supported: + raise UnsupportedError( + f"backend {self.backend!r} does not support {op.value}" + ) + cuda_ordering = _coerce_enum(cuda_ordering, CudaOrderingMode, "cuda_ordering") + _validate_cuda_ordering_arguments( + self.capabilities, + self.backend, + cuda_ordering, + stream, + ready_event, + ) + local_indices, remote_indices = _validate_plan_indices( + plan, local_indices, remote_indices, op + ) + notification = _validate_attached_notification( + self.capabilities, self.backend, notification + ) + cuda_device = _validate_cuda_ordering_selection( + plan, + local_indices, + cuda_ordering, + stream, + ready_event, + ) + binding = BoundTransfer( + self, + plan, + op, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + cuda_device=cuda_device, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + selection=None, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._bound_transfers.add(binding) + plan._bindings.add(binding) + except BaseException: + plan._bindings.discard(binding) + self._bound_transfers.discard(binding) + binding._closed = True + binding._plan = None + raise + return binding + + @_state_locked + def _bind_index_selection( + self, + selection: IndexSelection, + *, + cuda_ordering: CudaOrderingMode | str, + stream: object | None, + ready_event: object | None, + ) -> BoundTransfer: + self._check_open() + if type(selection) is not IndexSelection: + raise TypeError("expected an IndexSelection") + if selection._endpoint is not self: + raise InvalidRegionError("index selection belongs to another endpoint") + plan = selection._check_open() + cuda_ordering = _coerce_enum(cuda_ordering, CudaOrderingMode, "cuda_ordering") + _validate_cuda_ordering_arguments( + self.capabilities, + self.backend, + cuda_ordering, + stream, + ready_event, + ) + cuda_device = _validate_cuda_ordering_devices( + selection._local_devices, + cuda_ordering, + stream, + ready_event, + ) + binding = BoundTransfer( + self, + plan, + selection._op, + local_indices=selection._backend_selection.local_indices, + remote_indices=selection._backend_selection.remote_indices, + notification=None, + cuda_device=cuda_device, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + selection=selection, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._bound_transfers.add(binding) + plan._bindings.add(binding) + except BaseException: + plan._bindings.discard(binding) + self._bound_transfers.discard(binding) + binding._closed = True + binding._plan = None + binding._selection = None + raise + return binding + + @_state_locked + def _create_transfer_slot( + self, + bound: BoundTransfer, + *, + request_local_notifications: bool, + ) -> TransferSlot | RequestLocalNotificationSlot: + """Publish Core ownership before asking a provider for native state.""" + + self._check_open() + if not self._capabilities.persistent_request_slots: + raise UnsupportedError( + f"backend {self.backend!r} does not support persistent request slots" + ) + if not isinstance(bound, BoundTransfer): + raise TypeError("expected a BoundTransfer") + if bound._endpoint is not self: + raise InvalidRegionError("bound transfer belongs to another endpoint") + if request_local_notifications: + if not self._capabilities.request_slot_notification_overrides: + raise UnsupportedError( + f"backend {self.backend!r} does not support request-local " + "persistent-slot notifications" + ) + if bound._notification is not None: + raise ValueError( + "request-local notification slots require a binding with " + "notification=None" + ) + plan = bound._check_open() + slot: TransferSlot + if request_local_notifications: + slot = RequestLocalNotificationSlot(self, bound, _token=_CONSTRUCTION_TOKEN) + else: + slot = TransferSlot(self, bound, _token=_CONSTRUCTION_TOKEN) + slot._creation_failed = True + try: + # Publish the recovery registry first. At every subsequent + # interruption point either normal rollback succeeds or the slot + # remains discoverable through ``open_transfer_slots``. + self._transfer_slots.add(slot) + bound._slots.add(slot) + result = self._backend_create_request_slot( + bound._op, + plan._handle, + local_indices=bound._local_indices, + remote_indices=bound._remote_indices, + notification=bound._notification, + adopt_slot=slot._adopt_backend_slot, + ) + if result is not None: + raise BackendFailureError( + "backend request-slot creation must return None after adoption" + ) + if slot._backend_slot is None: + raise BackendFailureError( + "backend request-slot creation returned without adoption" + ) + slot._creation_failed = False + return slot + except BaseException as error: + if slot._backend_slot is not None: + try: + self._close_transfer_slot(slot) + except BaseException as cleanup_error: + slot._creation_failed = True + raise _creation_cleanup_failure( + "transfer slot", error, cleanup_error + ) from cleanup_error + else: + # Keep the endpoint recovery edge until the bound blocker has + # been removed. If cleanup itself is interrupted, the slot is + # still public through the endpoint snapshot and closeable. + bound._slots.discard(slot) + slot._closed = True + self._finish_transfer_slot_close(slot) + raise + + def _reserve_indexed_selection_submission( + self, + *, + plan: TransferPlan | None, + op: TransferOp | str | None, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + selection: IndexSelection | None, + notification: bytes | bytearray | memoryview | None, + cuda_ordering: CudaOrderingMode | str, + stream: object | None, + ready_event: object | None, + recovery_token: object | None, + ) -> tuple[ + TransferPlan, + TransferOp, + BackendIndexSelection, + bytes | None, + CudaOrderingMode, + object | None, + Work, + ]: + """Validate and construct an unpublished indexed Work under state lock. + + The caller holds the lock through publication. Returning a published + Work through a decorated helper would expose a lock-exit ownership gap. + """ + + active_plan: TransferPlan | None = None + work: Work | None = None + try: + self._check_open() + if selection is None: + if plan is None: + raise RuntimeError( + "one-shot indexed submission is missing its internal plan" + ) + self._validate_plan_locked(plan) + if not plan.indexed: + raise InvalidRegionError( + "indices require a plan prepared with indexed=True" + ) + active_plan = plan + active_op = _coerce_enum(op, TransferOp, "op") + supported = ( + self.capabilities.read + if active_op is TransferOp.READ + else self.capabilities.write + ) + if not supported: + raise UnsupportedError( + f"backend {self.backend!r} does not support {active_op.value}" + ) + caller_ready_hint = bool( + cuda_ordering is CudaOrderingMode.CALLER_READY + or ( + type(cuda_ordering) is str + and cuda_ordering == CudaOrderingMode.CALLER_READY.value + ) + ) + backend_selection, local_devices = _snapshot_backend_index_selection( + active_plan, + active_op, + local_indices, + remote_indices, + collect_local_devices=not caller_ready_hint, + ) + selection_dependency: object = backend_selection + else: + if type(selection) is not IndexSelection: + raise TypeError("expected an IndexSelection") + if selection._endpoint is not self: + raise InvalidRegionError( + "index selection belongs to another endpoint" + ) + active_plan = selection._plan + self._validate_plan_locked(active_plan) + active_op = selection._op + backend_selection = selection._backend_selection + local_devices = selection._local_devices + selection_dependency = selection + cuda_ordering = _coerce_enum( + cuda_ordering, CudaOrderingMode, "cuda_ordering" + ) + _validate_cuda_ordering_arguments( + self.capabilities, + self.backend, + cuda_ordering, + stream, + ready_event, + ) + notification = _validate_attached_notification( + self.capabilities, self.backend, notification + ) + cuda_device = ( + None + if cuda_ordering is CudaOrderingMode.CALLER_READY + else _validate_cuda_ordering_devices( + local_devices, + cuda_ordering, + stream, + ready_event, + ) + ) + if stream is None and ready_event is None: + dependencies = ( + (selection_dependency,) + if notification is None + else (selection_dependency, notification) + ) + else: + dependencies = tuple( + value + for value in ( + selection_dependency, + stream, + ready_event, + notification, + ) + if value is not None + ) + work = Work( + self, + None, + plan=active_plan, + dependencies=dependencies, + recovery_token=recovery_token, + _token=_CONSTRUCTION_TOKEN, + ) + return ( + active_plan, + active_op, + backend_selection, + notification, + cuda_ordering, + cuda_device, + work, + ) + except BaseException: + if work is not None: + work_lock = work._lock + if work_lock is None: + work._closed = True + work._submitting = False + work._dependencies = () + work._plan = None + work._recovery_token = None + else: + with work_lock: + work._closed = True + work._submitting = False + work._dependencies = () + work._plan = None + work._recovery_token = None + raise + + def _submit_indexed_selection( + self, + *, + plan: TransferPlan | None, + op: TransferOp | str | None, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + selection: IndexSelection | None, + notification: bytes | bytearray | memoryview | None, + cuda_ordering: CudaOrderingMode | str, + stream: object | None, + ready_event: object | None, + recovery_token: object | None, + initial_max_polls: int | None = None, + initial_timeout_ns: int | None = None, + ) -> Work: + """Submit a reusable public or one-shot private indexed selection.""" + + active_plan: TransferPlan | None = None + active_op: TransferOp | None = None + backend_selection: BackendIndexSelection | None = None + cuda_device: object | None = None + submission_token: object | None = None + work: Work | None = None + try: + state_lock = self._state_lock + if state_lock is _NOOP_LOCK: + ( + active_plan, + active_op, + backend_selection, + notification, + cuda_ordering, + cuda_device, + work, + ) = self._reserve_indexed_selection_submission( + plan=plan, + op=op, + local_indices=local_indices, + remote_indices=remote_indices, + selection=selection, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + submission_token = work + self._works.add(work) + active_plan._submitting.add(work) + active_plan._works.add(work) + else: + with state_lock: + ( + active_plan, + active_op, + backend_selection, + notification, + cuda_ordering, + cuda_device, + work, + ) = self._reserve_indexed_selection_submission( + plan=plan, + op=op, + local_indices=local_indices, + remote_indices=remote_indices, + selection=selection, + notification=notification, + cuda_ordering=cuda_ordering, + stream=stream, + ready_event=ready_event, + recovery_token=recovery_token, + ) + submission_token = work + self._works.add(work) + active_plan._submitting.add(work) + active_plan._works.add(work) + + if cuda_ordering is not CudaOrderingMode.CALLER_READY: + _establish_cuda_ordering_prevalidated( + cuda_ordering, + stream, + ready_event, + cuda_device, + ) + if initial_max_polls is None: + result = self._backend_submit_prevalidated( + active_op, + active_plan._handle, + selection=backend_selection, + notification=notification, + # A direct selection submit is the allocation-minimal + # dynamic form used once by framework requests. Bind + # explicitly for native reuse across generations. + reuse_token=None, + adopt_work=work._adopt_backend_work, + ) + if result is not None: + raise BackendFailureError( + "backend prevalidated submit must return None after " + "adopting work" + ) + else: + result = self._backend_submit_prevalidated_and_poll( + active_op, + active_plan._handle, + selection=backend_selection, + notification=notification, + max_polls=initial_max_polls, + timeout_ns=initial_timeout_ns, + adopt_work=work._adopt_backend_work, + ) + current_work = work + work_lock = current_work._lock + if work_lock is None: + if current_work._backend_work is None: + raise BackendFailureError( + "backend prevalidated submit returned without adopting work" + ) + if initial_max_polls is not None: + current_work._ingest_backend_state_unlocked(result) + else: + with work_lock: + if current_work._backend_work is None: + raise BackendFailureError( + "backend prevalidated submit returned without adopting work" + ) + if initial_max_polls is not None: + current_work._ingest_backend_state_unlocked(result) + if state_lock is _NOOP_LOCK: + active_plan._submitting.discard(submission_token) + current_work._submitting = False + else: + with state_lock: + active_plan._submitting.discard(submission_token) + with current_work._lock or _NOOP_LOCK: + current_work._submitting = False + return current_work + except BaseException as error: + state_lock = self._state_lock + if state_lock is _NOOP_LOCK: + adopted = self._reconcile_failed_indexed_submission( + active_plan, work, submission_token, error + ) + else: + with state_lock: + adopted = self._reconcile_failed_indexed_submission( + active_plan, work, submission_token, error + ) + if adopted: + return cast(Work, work) + raise + + def _reconcile_failed_indexed_submission( + self, + active_plan: TransferPlan | None, + work: Work | None, + submission_token: object | None, + error: BaseException, + ) -> bool: + """Repair indexed submission while the caller holds state lock.""" + + if active_plan is not None and work is not None: + active_plan._submitting.discard( + work if submission_token is None else submission_token + ) + if work is None: + return False + work_lock = work._lock + if work_lock is None: + return self._reconcile_failed_indexed_work_unlocked( + active_plan, work, error + ) + with work_lock: + return self._reconcile_failed_indexed_work_unlocked( + active_plan, work, error + ) + + def _reconcile_failed_indexed_work_unlocked( + self, + active_plan: TransferPlan | None, + work: Work, + error: BaseException, + ) -> bool: + if work._backend_work is not None: + work._submission_error = error + if work._state in _TERMINAL_STATES: + work._state = WorkState.FAILED + work._error = error + work._submitting = False + return True + if active_plan is not None: + active_plan._works.discard(work) + self._works.discard(work) + work._closed = True + work._submitting = False + work._dependencies = () + work._plan = None + work._recovery_token = None + return False + + def _submit_bound( + self, + bound: BoundTransfer, + *, + notification: object = _UNSET_NOTIFICATION, + recovery_token: object | None, + ) -> Work: + """Submit a binding without repeating layout or storage validation.""" + + plan: TransferPlan | None = None + selection: IndexSelection | None = None + active_notification: bytes | None = None + reuse_token: object | None = None + work: Work | None = None + try: + state_lock = self._state_lock + with state_lock: + self._check_open() + if not isinstance(bound, BoundTransfer): + raise TypeError("expected a BoundTransfer") + if bound._endpoint is not self: + raise InvalidRegionError( + "bound transfer belongs to another endpoint" + ) + if bound._closed or bound._plan is None: + raise ClosedError("bound transfer is closed") + plan = bound._plan + if plan._closed: + raise ClosedError("transfer plan is closed") + if plan._creation_failed: + raise BackendFailureError( + "transfer plan is retained for cleanup only" + ) + selection = bound._selection + if selection is None: + if notification is not _UNSET_NOTIFICATION: + raise TypeError( + "notification override requires an IndexSelection binding" + ) + active_notification = bound._notification + else: + active_notification = ( + None + if notification is _UNSET_NOTIFICATION + else _validate_attached_notification( + self.capabilities, + self.backend, + notification, + ) + ) + reuse_token = selection._reuse_token + if reuse_token is None: + reuse_token = _ReuseToken() + object.__setattr__(selection, "_reuse_token", reuse_token) + work = Work( + self, + None, + plan=plan, + bound=bound, + dependencies=( + () if active_notification is None else (active_notification,) + ), + recovery_token=recovery_token, + _token=_CONSTRUCTION_TOKEN, + ) + # Use Work itself as the in-progress token. This publishes all + # lifetime reservations before CUDA or provider activity while + # avoiding a second per-submission allocation. + self._works.add(work) + plan._submitting.add(work) + bound._submitting.add(work) + plan._works.add(work) + bound._works.add(work) + + if bound._cuda_ordering is not CudaOrderingMode.CALLER_READY: + _establish_cuda_ordering_prevalidated( + bound._cuda_ordering, + bound._stream, + bound._ready_event, + bound._cuda_device, + ) + + if selection is None: + result = self._backend_submit( + bound._op, + plan._handle, + local_indices=bound._local_indices, + remote_indices=bound._remote_indices, + notification=active_notification, + reuse_token=bound._reuse_token, + adopt_work=work._adopt_backend_work, + ) + else: + result = self._backend_submit_prevalidated( + bound._op, + plan._handle, + selection=selection._backend_selection, + notification=active_notification, + reuse_token=reuse_token, + adopt_work=work._adopt_backend_work, + ) + if result is not None: + raise BackendFailureError( + "backend submit must return None after adopting its work" + ) + current_work = work + work_lock = current_work._lock + if work_lock is None: + if current_work._backend_work is None: + raise BackendFailureError( + "backend submit returned without adopting transfer work" + ) + else: + with work_lock: + if current_work._backend_work is None: + raise BackendFailureError( + "backend submit returned without adopting transfer work" + ) + if state_lock is _NOOP_LOCK: + plan._submitting.discard(current_work) + bound._submitting.discard(current_work) + current_work._submitting = False + else: + with state_lock: + plan._submitting.discard(current_work) + bound._submitting.discard(current_work) + with current_work._lock or _NOOP_LOCK: + current_work._submitting = False + return current_work + except BaseException as error: + with self._state_lock: + if plan is not None and work is not None: + plan._submitting.discard(work) + bound._submitting.discard(work) + adopted = False + if work is not None: + with work._lock or _NOOP_LOCK: + adopted = work._backend_work is not None + if adopted: + # Preserve the generic submission recovery contract: once + # adopted, an interruption is reported by the retained Work. + work = cast(Work, work) + with work._lock or _NOOP_LOCK: + work._submission_error = error + work._submitting = False + return work + if work is not None: + if plan is not None: + plan._works.discard(work) + bound._works.discard(work) + self._works.discard(work) + with work._lock or _NOOP_LOCK: + work._closed = True + work._submitting = False + work._dependencies = () + work._plan = None + work._bound = None + work._recovery_token = None + raise + + def send_notification( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + *, + recovery_token: object | None = None, + ) -> Work: + if self._thread_mode is ThreadMode.MULTIPLE: + return self._send_notification_multiple( + peer, payload, recovery_token=recovery_token + ) + + state_lock = self._state_lock + if state_lock is _NOOP_LOCK: + return self._send_notification_locked( + peer, payload, recovery_token=recovery_token + ) + with state_lock: + return self._send_notification_locked( + peer, payload, recovery_token=recovery_token + ) + + def _send_notification_locked( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + *, + recovery_token: object | None, + ) -> Work: + """Send while the selected non-MULTIPLE lifecycle boundary is held.""" + + self._check_open() + self._validate_peer(peer) + if not self.capabilities.standalone_notifications: + raise UnsupportedError( + f"backend {self.backend!r} does not support notifications" + ) + payload = _as_bytes(payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES) + work = Work( + self, + None, + dependencies=(peer, payload), + recovery_token=recovery_token, + _token=_CONSTRUCTION_TOKEN, + ) + try: + self._works.add(work) + result = self._backend_send_notification( + peer._handle, + payload, + adopt_work=work._adopt_backend_work, + ) + if result is not None: + raise BackendFailureError( + "backend notification send must return None after adopting its work" + ) + lock = work._lock + if lock is None: + if work._backend_work is None: + raise BackendFailureError( + "backend notification send returned without adopting work" + ) + else: + with lock: + if work._backend_work is None: + raise BackendFailureError( + "backend notification send returned without adopting work" + ) + work._submitting = False + return work + except BaseException as error: + with work._lock or _NOOP_LOCK: + adopted = work._backend_work is not None + if adopted: + work._submission_error = error + work._submitting = False + return work + work._closed = True + work._submitting = False + work._dependencies = () + work._recovery_token = None + self._works.discard(work) + raise + + def _send_notification_multiple( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + *, + recovery_token: object | None, + ) -> Work: + """Publish Work ownership, then preserve provider MULTIPLE concurrency.""" + + work: Work | None = None + try: + with self._state_lock: + self._check_open() + self._validate_peer(peer) + if not self._capabilities.standalone_notifications: + raise UnsupportedError( + f"backend {self._backend_name!r} does not support notifications" + ) + payload = _as_bytes( + payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES + ) + work = Work( + self, + None, + dependencies=(peer, payload), + recovery_token=recovery_token, + _token=_CONSTRUCTION_TOKEN, + ) + # Work itself is the allocation already required by this API and + # becomes the close-visible lifetime reservation before native + # entry. No second per-send token is needed. + self._works.add(work) + + result = self._backend_send_notification( + peer._handle, + payload, + adopt_work=work._adopt_backend_work, + ) + if result is not None: + raise BackendFailureError( + "backend notification send must return None after adopting its work" + ) + with work._lock or _NOOP_LOCK: + if work._backend_work is None: + raise BackendFailureError( + "backend notification send returned without adopting work" + ) + with self._state_lock: + with work._lock or _NOOP_LOCK: + work._submitting = False + return work + except BaseException as error: + with self._state_lock: + if work is None: + raise + with work._lock or _NOOP_LOCK: + if work._backend_work is not None: + work._submission_error = error + work._submitting = False + return work + work._closed = True + work._submitting = False + work._dependencies = () + work._recovery_token = None + self._works.discard(work) + raise + + def send_notification_sync( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + ) -> None: + """Send without allocating a :class:`Work`. + + Return means exactly ``standalone_notification_completion``. In + particular, NIXL's ``LOCAL_ACCEPTED`` does not mean the peer observed + the payload. Once backend entry begins, an exception can be + outcome-ambiguous; protocols that need certainty must use sequence + numbers and acknowledgement, or retain :meth:`send_notification` and + its inspectable Work ownership. + """ + + if self._thread_mode is ThreadMode.MULTIPLE: + self._send_notification_sync_multiple(peer, payload) + return + + state_lock = self._state_lock + if state_lock is _NOOP_LOCK: + self._send_notification_sync_locked(peer, payload) + return + with state_lock: + self._send_notification_sync_locked(peer, payload) + + def _send_notification_sync_locked( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + ) -> None: + """Validate and send while the caller supplies lifecycle exclusion.""" + + self._check_open() + self._validate_peer(peer) + if not self._capabilities.synchronous_notification_send: + raise UnsupportedError( + f"backend {self._backend_name!r} does not support synchronous " + "notification send" + ) + payload = _as_bytes(payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES) + send = self._send_notification_sync_dispatch + if send is None: + raise BackendFailureError( + "backend advertised synchronous notification send without a " + "dispatch method" + ) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = send(peer._handle, payload) + else: + result = send(peer._handle, payload) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._backend_name!r} failed to send notification " + "synchronously; delivery outcome may be ambiguous" + ) from error + if result is not None: + raise BackendFailureError( + "backend synchronous notification send must return None; " + "delivery outcome may be ambiguous" + ) + + def _send_notification_sync_multiple( + self, + peer: Peer, + payload: bytes | bytearray | memoryview, + ) -> None: + """Reserve peer lifetime briefly, then preserve native RW concurrency.""" + + reserved_peer: Peer | None = None + reservation: _SyncNotificationReservation | None = None + try: + with self._state_lock: + self._check_open() + self._validate_peer(peer) + if not self._capabilities.synchronous_notification_send: + raise UnsupportedError( + f"backend {self._backend_name!r} does not support synchronous " + "notification send" + ) + payload = _as_bytes( + payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES + ) + send = self._send_notification_sync_dispatch + if send is None: + raise BackendFailureError( + "backend advertised synchronous notification send without a " + "dispatch method" + ) + peer_handle = peer._handle + pool = peer._sync_notification_reservation_pool + reservation = pool.pop() if pool else _SyncNotificationReservation() + reserved_peer = peer + reservation.active = True + peer._active_sync_notification_sends.add(reservation.reference) + + try: + result = send(peer_handle, payload) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self._backend_name!r} failed to send notification " + "synchronously; delivery outcome may be ambiguous" + ) from error + if result is not None: + raise BackendFailureError( + "backend synchronous notification send must return None; " + "delivery outcome may be ambiguous" + ) + finally: + if reserved_peer is not None and reservation is not None: + # Publish native exit before taking an interruptible cleanup + # lock. A retained exception traceback may keep this frame and + # token alive, so liveness of the weakref alone is insufficient. + reservation.active = False + with self._state_lock: + reserved_peer._active_sync_notification_sends.discard( + reservation.reference + ) + if not reserved_peer._closed: + reserved_peer._sync_notification_reservation_pool.append( + reservation + ) + + def poll_notifications(self, max_items: int | None = None) -> list[Notification]: + if self._thread_mode is ThreadMode.MULTIPLE: + notifications = self._poll_backend_notifications_multiple( + max_items, grouped=False + ) + else: + state_lock = self._state_lock + with state_lock: + self._check_open() + if not self.capabilities.notifications: + raise UnsupportedError( + f"backend {self.backend!r} does not support notifications" + ) + if max_items is not None: + if isinstance(max_items, bool) or not isinstance(max_items, int): + raise TypeError( + "max_items must be a non-negative integer or None" + ) + if max_items < 0: + raise ValueError("max_items must be non-negative") + if max_items == 0: + return [] + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + notifications = self._backend.poll_notifications(max_items) + else: + notifications = self._backend.poll_notifications(max_items) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to poll notifications" + ) from error + if not isinstance(notifications, list) or not all( + isinstance(item, BackendNotification) for item in notifications + ): + raise BackendFailureError( + "backend notifications must be a list of BackendNotification" + ) + if max_items is not None and len(notifications) > max_items: + raise BackendFailureError( + "backend returned more notifications than max_items" + ) + return [ + Notification( + payload=notification.payload, + source_endpoint_id=notification.source_endpoint_id, + source_incarnation=notification.source_incarnation, + ) + for notification in notifications + ] + + def poll_notification_batches( + self, max_items: int | None = None + ) -> list[NotificationBatch]: + """Destructively receive source-grouped notifications. + + ``max_items`` counts payloads rather than returned batches. The v0 + receive contract remains at-most-once: there is no acknowledgement or + rollback token. Per-source payload order is preserved, but providers + need not define a total order across different sources and may return + repeated batches for one source. + """ + + if self._thread_mode is ThreadMode.MULTIPLE: + batches = self._poll_backend_notifications_multiple(max_items, grouped=True) + else: + state_lock = self._state_lock + with state_lock: + self._check_open() + if not self.capabilities.notifications: + raise UnsupportedError( + f"backend {self.backend!r} does not support notifications" + ) + if max_items is not None: + if isinstance(max_items, bool) or not isinstance(max_items, int): + raise TypeError( + "max_items must be a non-negative integer or None" + ) + if max_items < 0: + raise ValueError("max_items must be non-negative") + if max_items == 0: + return [] + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + batches = self._backend.poll_notification_batches(max_items) + else: + batches = self._backend.poll_notification_batches(max_items) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to poll notification batches" + ) from error + if not isinstance(batches, list) or not all( + type(item) is BackendNotificationBatch for item in batches + ): + raise BackendFailureError( + "backend notification batches must be a list of " + "BackendNotificationBatch" + ) + if max_items is not None: + payload_count = sum(len(batch.payloads) for batch in batches) + if payload_count > max_items: + raise BackendFailureError( + "backend returned more notification payloads than max_items" + ) + return [_notification_batch_from_backend(batch) for batch in batches] + + def _poll_backend_notifications_multiple( + self, + max_items: int | None, + *, + grouped: bool, + ) -> object: + """Run one destructive receive with narrow, close-visible admission.""" + + reservation: _BackendCallReservation | None = None + # Serialize only destructive consumers. Acquire this lock before call + # admission so queued consumers need no reservation and can safely + # discover a concurrently completed close before backend entry. + with self._notification_consumer_lock: + try: + with self._state_lock: + self._check_open() + if not self._capabilities.notifications: + raise UnsupportedError( + f"backend {self._backend_name!r} does not support " + "notifications" + ) + if max_items is not None: + if isinstance(max_items, bool) or not isinstance( + max_items, int + ): + raise TypeError( + "max_items must be a non-negative integer or None" + ) + if max_items < 0: + raise ValueError("max_items must be non-negative") + if max_items == 0: + return [] + poll = ( + self._backend.poll_notification_batches + if grouped + else self._backend.poll_notifications + ) + reservation = self._reserve_multiple_backend_call_locked() + try: + return poll(max_items) + except TransferError: + raise + except Exception as error: + action = ( + "poll notification batches" if grouped else "poll notifications" + ) + raise BackendFailureError( + f"backend {self._backend_name!r} failed to {action}" + ) from error + finally: + if reservation is not None: + self._finish_multiple_backend_call(reservation) + + def progress(self) -> None: + if self._thread_mode is ThreadMode.MULTIPLE: + self._progress_multiple() + return + + state_lock = self._state_lock + if state_lock is _NOOP_LOCK: + self._progress_locked() + return + with state_lock: + self._progress_locked() + + def _progress_locked(self) -> None: + self._check_open() + if not self.capabilities.manual_progress: + raise UnsupportedError( + f"backend {self.backend!r} does not support manual progress" + ) + self._backend_call("progress backend", self._backend.progress) + + def _progress_multiple(self) -> None: + reservation: _BackendCallReservation | None = None + try: + with self._state_lock: + self._check_open() + if not self._capabilities.manual_progress: + raise UnsupportedError( + f"backend {self._backend_name!r} does not support manual progress" + ) + progress = self._backend.progress + reservation = self._reserve_multiple_backend_call_locked() + self._backend_call("progress backend", progress) + finally: + if reservation is not None: + self._finish_multiple_backend_call(reservation) + + @_state_locked + def close(self) -> None: + if self._creation_failed: + self._close_failed_endpoint_backend() + return + if self._closed: + return + if not self._closing: + if self._multiple_backend_calls_active_locked(): + raise BusyError( + "cannot close endpoint while backend calls are in progress" + ) + self._retry_orphaned_plans() + active = [work for work in self._works if not work.test()] + if active: + raise BusyError( + f"cannot close endpoint with {len(active)} active work item(s)" + ) + counts = { + "work": len(self._works), + "transfer slot": len(self._transfer_slots), + "bound transfer": len(self._bound_transfers), + "plan": len(self._plans), + "peer": len(self._peers), + "registration": len(self._registrations), + } + if any(counts.values()): + detail = ", ".join( + f"{name}={count}" for name, count in counts.items() if count + ) + raise BusyError( + f"close child resources before endpoint close ({detail})" + ) + # From this point only close() may be retried. This makes a signal + # after successful provider teardown converge without calling + # progress on an already closed provider. + self._closing = True + self._backend_call( + "close backend", + lambda: self._backend_receipt.close(self._backend.close), + ) + self._detach_backend_finalizer() + self._closed = True + active_calls = self._active_multiple_backend_calls + if active_calls is not None: + active_calls.clear() + reservation_pool = self._multiple_backend_call_reservation_pool + if reservation_pool is not None: + reservation_pool.clear() + + def _close_failed_endpoint_backend(self) -> None: + """Retry backend-only cleanup after Endpoint initialization failed.""" + + if self._closed: + self._creation_failed = False + self._detach_backend_finalizer() + return + self._backend_call( + "close backend after failed endpoint initialization", + lambda: self._backend_receipt.close(self._backend.close), + ) + self._detach_backend_finalizer() + self._closed = True + self._creation_failed = False + + def _cleanup_adopted_backend_after_creation_error( + self, original_error: BaseException + ) -> None: + implementation = self._backend_receipt.backend + if implementation is None: + self._detach_backend_finalizer() + return + self._backend = implementation + self._creation_failed = True + try: + self._close_failed_endpoint_backend() + except BaseException as cleanup_error: + raise _endpoint_creation_cleanup_failure( + self, original_error, cleanup_error + ) from cleanup_error + + def _detach_backend_finalizer(self) -> None: + self._backend_finalizer.detach() + + def _submit( + self, + op: TransferOp, + plan: TransferPlan | None, + *, + local: Sequence[Region] | None, + remote: Sequence[Region] | None, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | bytearray | memoryview | None, + cuda_ordering: CudaOrderingMode | str, + stream: object | None, + ready_event: object | None, + recovery_token: object | None, + ) -> Work: + owned_plan = False + plan_to_release: TransferPlan | None = None + reserved_plan: TransferPlan | None = None + submission_token: object | None = None + work: Work | None = None + try: + state_lock = self._state_lock + with state_lock: + self._check_open() + supported = ( + self.capabilities.read + if op is TransferOp.READ + else self.capabilities.write + ) + if not supported: + raise UnsupportedError( + f"backend {self.backend!r} does not support {op.value}" + ) + cuda_ordering = _coerce_enum( + cuda_ordering, CudaOrderingMode, "cuda_ordering" + ) + _validate_cuda_ordering_arguments( + self.capabilities, + self.backend, + cuda_ordering, + stream, + ready_event, + ) + if plan is None: + if local is None or remote is None: + raise TypeError( + "one-shot transfer requires both local and remote regions" + ) + plan = self.prepare(local=local, remote=remote, indexed=False) + plan_to_release = plan + owned_plan = True + elif local is not None or remote is not None: + raise TypeError("pass a plan or one-shot regions, not both") + active_plan = plan + self._validate_plan_locked(active_plan) + local_indices, remote_indices = _validate_plan_indices( + active_plan, local_indices, remote_indices, op + ) + notification = _validate_attached_notification( + self.capabilities, self.backend, notification + ) + cuda_device = _validate_cuda_ordering_selection( + active_plan, + local_indices, + cuda_ordering, + stream, + ready_event, + ) + dependencies = tuple( + value + for value in ( + local_indices, + remote_indices, + stream, + ready_event, + notification, + ) + if value is not None + ) + # Publish the complete Core lifetime graph before entering the + # provider. The provider adopts its request into this wrapper + # before it is allowed to launch transport work. + work = Work( + self, + None, + plan=active_plan, + dependencies=dependencies, + owned_plan=owned_plan, + recovery_token=recovery_token, + _token=_CONSTRUCTION_TOKEN, + ) + submission_token = work + reserved_plan = active_plan + self._works.add(work) + active_plan._submitting.add(submission_token) + active_plan._works.add(work) + _establish_cuda_ordering_prevalidated( + cuda_ordering, + stream, + ready_event, + cuda_device, + ) + + reuse_token = ( + active_plan._reuse_tokens[op] + if not owned_plan + and local_indices is None + and remote_indices is None + and notification is None + and cuda_ordering is CudaOrderingMode.CALLER_READY + else None + ) + result = self._backend_submit( + op, + active_plan._handle, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + reuse_token=reuse_token, + adopt_work=work._adopt_backend_work, + ) + if result is not None: + raise BackendFailureError( + "backend submit must return None after adopting its work" + ) + current_work = work + work_lock = current_work._lock + if work_lock is None: + if current_work._backend_work is None: + raise BackendFailureError( + "backend submit returned without adopting transfer work" + ) + else: + with work_lock: + if current_work._backend_work is None: + raise BackendFailureError( + "backend submit returned without adopting transfer work" + ) + if state_lock is _NOOP_LOCK: + reserved_plan._submitting.discard(submission_token) + current_work._submitting = False + else: + with state_lock: + reserved_plan._submitting.discard(submission_token) + with current_work._lock or _NOOP_LOCK: + current_work._submitting = False + return current_work + except BaseException as error: + with self._state_lock: + if reserved_plan is not None and submission_token is not None: + reserved_plan._submitting.discard(submission_token) + adopted = False + if work is not None: + with work._lock or _NOOP_LOCK: + adopted = work._backend_work is not None + if adopted: + # A provider must not raise after adoption, but retaining + # the Work is the only safe response if it does: native + # access may already be in flight. Surface the protocol + # violation only after the backend request is terminal. + work = cast(Work, work) + with work._lock or _NOOP_LOCK: + work._submission_error = error + work._submitting = False + return work + if work is not None: + if work._plan is not None: + work._plan._works.discard(work) + self._works.discard(work) + with work._lock or _NOOP_LOCK: + work._closed = True + work._submitting = False + work._dependencies = () + work._plan = None + work._recovery_token = None + if plan_to_release is not None: + try: + plan_to_release.close() + except BaseException: + self._orphaned_plans.add(plan_to_release) + raise + + def _validate_regions( + self, + local: tuple[Region, ...], + remote: tuple[Region, ...], + *, + indexed: bool, + ) -> None: + if not local or not remote: + raise InvalidRegionError("local and remote regions must be non-empty") + for region in local: + if not isinstance(region, Region) or region.is_remote: + raise InvalidRegionError( + "local must contain local Registration regions" + ) + if not isinstance(region._origin, Registration): + raise InvalidRegionError("local region has an invalid owner") + self._validate_registration(region._origin) + peers: set[Peer] = set() + for region in remote: + if not isinstance(region, Region) or not region.is_remote: + raise InvalidRegionError("remote must contain Peer regions") + if not isinstance(region._origin, Peer): + raise InvalidRegionError("remote region has an invalid owner") + self._validate_peer(region._origin) + peers.add(region._origin) + if len(peers) != 1: + raise InvalidRegionError("one transfer plan may target exactly one peer") + if not indexed: + local_size = sum(region.total_nbytes for region in local) + remote_size = sum(region.total_nbytes for region in remote) + if local_size != remote_size: + raise InvalidRegionError( + f"local and remote byte totals differ: " + f"{local_size} != {remote_size}" + ) + + def _validate_registration(self, registration: Registration) -> None: + if not isinstance(registration, Registration): + raise TypeError("expected a Registration") + if registration._endpoint is not self: + raise InvalidRegionError("registration belongs to another endpoint") + registration._check_open() + _validate_memory_stable(registration._memory) + + def _validate_peer(self, peer: Peer) -> None: + if not isinstance(peer, Peer): + raise TypeError("expected a Peer") + if peer._endpoint is not self: + raise InvalidRegionError("peer belongs to another endpoint") + peer._check_open() + + def _validate_plan(self, plan: TransferPlan) -> None: + self._check_open() + self._validate_plan_locked(plan) + + def _validate_plan_locked(self, plan: TransferPlan) -> None: + """Validate a plan after this endpoint has entered its state lock.""" + + if not isinstance(plan, TransferPlan): + raise TypeError("expected a TransferPlan") + if plan._endpoint is not self: + raise InvalidRegionError("transfer plan belongs to another endpoint") + if plan._closed: + raise ClosedError("transfer plan is closed") + if plan._creation_failed: + raise BackendFailureError("transfer plan is retained for cleanup only") + for origin in plan._origins: + if isinstance(origin, Registration): + if origin._endpoint is not self: + raise InvalidRegionError("registration belongs to another endpoint") + if origin._closed: + raise ClosedError("registration is closed") + if origin._creation_failed: + raise BackendFailureError( + "registration is retained for cleanup only" + ) + _validate_memory_stable(origin._memory) + + @_state_locked + def _close_registration(self, registration: Registration) -> None: + if registration._closed: + self._finish_registration_close(registration) + return + self._check_open() + self._retry_orphaned_plans() + if any( + not plan._closed and registration in plan._origins for plan in self._plans + ): + raise BusyError( + f"registration {registration.name!r} is used by an open plan" + ) + result = self._backend_call( + "deregister memory", + lambda: self._backend.deregister(registration._handle), + ) + if result is not None: + raise BackendFailureError("backend deregister must return None") + registration._closed = True + self._finish_registration_close(registration) + + def _finish_registration_close(self, registration: Registration) -> None: + """Idempotently finish Core-only registration teardown.""" + + registration_id = registration.registration_id + name = registration.name + if self._registration_names.get(name) is registration: + self._registration_names.pop(name, None) + registration._handle = None + registration._memory = _released_memory(registration._memory) + # Authoritative recovery registry removal is the Core teardown commit. + if self._registrations.get(registration_id) is registration: + self._registrations.pop(registration_id, None) + + @_state_locked + def _close_peer(self, peer: Peer) -> None: + if peer._closed: + self._finish_peer_close(peer) + return + self._check_open() + self._retry_orphaned_plans() + active_sends = peer._active_sync_notification_sends + for reference in tuple(active_sends): + reservation = reference() + if reservation is None or not reservation.active: + active_sends.discard(reference) + if active_sends: + raise BusyError( + f"peer {peer.name!r} has synchronous notification sends in progress" + ) + if any(not plan._closed and peer in plan._origins for plan in self._plans): + raise BusyError(f"peer {peer.name!r} is used by an open plan") + if any(peer in work._dependencies for work in self._works): + raise BusyError(f"peer {peer.name!r} is used by an open work item") + result = self._backend_call( + "release peer", lambda: self._backend.release_peer(peer._handle) + ) + if result is not None: + raise BackendFailureError("backend release_peer must return None") + peer._closed = True + self._finish_peer_close(peer) + + def _finish_peer_close(self, peer: Peer) -> None: + """Idempotently finish Core-only peer teardown.""" + + if self._peers.get(peer.endpoint_id) is peer: + self._peers.pop(peer.endpoint_id, None) + peer._handle = None + peer._registrations = {} + peer._active_sync_notification_sends.clear() + peer._sync_notification_reservation_pool.clear() + + @_state_locked + def _close_plan(self, plan: TransferPlan) -> None: + if plan._closed: + self._finish_plan_close(plan) + return + self._check_open() + if plan._submitting: + raise BusyError("transfer plan has submissions in progress") + if plan._works: + raise BusyError("transfer plan has open work items") + if plan._bindings: + raise BusyError("transfer plan has open bound transfers") + result = self._backend_call( + "release transfer plan", + lambda: self._backend.release_plan(plan._handle), + ) + if result is not None: + raise BackendFailureError("backend release_plan must return None") + plan._closed = True + self._finish_plan_close(plan) + + def _finish_plan_close(self, plan: TransferPlan) -> None: + """Idempotently finish Core-only transfer-plan teardown.""" + + self._plans.discard(plan) + self._orphaned_plans.discard(plan) + plan._handle = None + plan._local = () + plan._remote = () + plan._origins = frozenset() + plan._reuse_tokens = {} + + @_state_locked + def _close_bound_transfer(self, bound: BoundTransfer) -> None: + if bound._closed: + # Closure is the commit marker, not an early-return condition. An + # interrupt after it may leave a parent edge or retained field; + # retrying completes the idempotent Core-only epilogue. + self._finish_bound_transfer_close(bound) + return + self._check_open() + if bound._endpoint is not self: + raise InvalidRegionError("bound transfer belongs to another endpoint") + plan = bound._check_open() + if bound._slots: + raise BusyError("bound transfer has open persistent request slots") + if bound._works: + raise BusyError("bound transfer has open work items") + if bound._submitting: + raise BusyError("bound transfer has submissions in progress") + # Unlink blocking graph edges before publishing CLOSED, while retaining + # `_plan` so a retry after that commit can repair either edge. + plan._bindings.discard(bound) + self._bound_transfers.discard(bound) + bound._closed = True + self._finish_bound_transfer_close(bound) + + def _finish_bound_transfer_close(self, bound: BoundTransfer) -> None: + """Idempotently finish Core-only bound-transfer unlink.""" + + plan = bound._plan + if plan is not None: + plan._bindings.discard(bound) + self._bound_transfers.discard(bound) + bound._plan = None + bound._local_indices = None + bound._remote_indices = None + bound._notification = None + bound._cuda_device = None + bound._stream = None + bound._ready_event = None + bound._reuse_token = None + bound._selection = None + + @_state_locked + def _close_transfer_slot(self, slot: TransferSlot) -> None: + """Release one idle persistent request and unlink its lifetime graph.""" + + if not isinstance(slot, TransferSlot): + raise TypeError("expected a TransferSlot") + if slot._endpoint is not self: + raise InvalidRegionError("transfer slot belongs to another endpoint") + if slot._closed: + # Closure is the commit marker, not an early-return condition. An + # interrupt after it may still leave one parent edge or retained + # callback, so every retry finishes the idempotent Core epilogue. + self._finish_transfer_slot_close(slot) + return + self._check_open() + bound = slot._bound + if bound is None or bound._closed: + raise ClosedError("transfer slot's bound transfer is closed") + if slot._active: + raise BusyError( + "transfer slot owns an unrecycled generation; recycle it before close" + ) + close = slot._backend_close + if close is None: + if slot._backend_slot is not None: + raise BackendFailureError( + "transfer slot has no backend close operation" + ) + else: + self._backend_request_slot_close(close) + + # Closure commits provider release while both recovery edges still + # exist. The retry path can therefore remove the bound blocker first + # and the endpoint's public recovery root last. + slot._closed = True + self._finish_transfer_slot_close(slot) + + def _finish_transfer_slot_close(self, slot: TransferSlot) -> None: + """Idempotently finish Core-only unlink after provider release.""" + + bound = slot._bound + if bound is not None: + bound._slots.discard(slot) + self._transfer_slots.discard(slot) + slot._bound = None + slot._backend_slot = None + slot._backend_execution_dispatch = None + slot._backend_start = None + slot._backend_start_and_poll = None + if isinstance(slot, RequestLocalNotificationSlot): + slot._backend_start_with_notification = None + slot._backend_start_and_poll_with_notification = None + slot._notification_fast_execute = None + slot._backend_poll_bounded = None + slot._backend_poll_state = None + slot._backend_cancel = None + slot._backend_recycle = None + slot._backend_close = None + slot._ready_event_synchronize = None + slot._fast_execute = None + slot._backend_generation_active = False + slot._error = None + slot._start_error = None + + def _retry_orphaned_plans(self) -> None: + for plan in tuple(self._orphaned_plans): + plan.close() + + def _check_open(self) -> None: + self._check_thread() + if self._creation_failed: + raise BackendFailureError( + f"endpoint {self.name!r} is retained for creation cleanup only; " + "retry close()" + ) + if self._closed: + raise ClosedError(f"endpoint {self.name!r} is closed") + if self._closing: + raise ClosedError( + f"endpoint {self.name!r} is closing; only close() may be retried" + ) + + def _reserve_multiple_backend_call_locked(self) -> _BackendCallReservation: + active_calls = self._active_multiple_backend_calls + pool = self._multiple_backend_call_reservation_pool + assert active_calls is not None and pool is not None + reservation = pool.pop() if pool else _BackendCallReservation() + try: + reservation.active = True + active_calls.add(reservation.reference) + except BaseException: + reservation.active = False + active_calls.discard(reservation.reference) + pool.append(reservation) + raise + return reservation + + def _finish_multiple_backend_call( + self, reservation: _BackendCallReservation + ) -> None: + # Publish native exit before interruptible cleanup. A traceback can keep + # this frame alive; close therefore also checks the explicit active bit. + reservation.active = False + with self._state_lock: + active_calls = self._active_multiple_backend_calls + pool = self._multiple_backend_call_reservation_pool + assert active_calls is not None and pool is not None + active_calls.discard(reservation.reference) + if not self._closing and not self._closed: + pool.append(reservation) + + def _multiple_backend_calls_active_locked(self) -> bool: + active_calls = self._active_multiple_backend_calls + if not active_calls: + return False + for reference in tuple(active_calls): + reservation = reference() + if reservation is None or not reservation.active: + active_calls.discard(reference) + return bool(active_calls) + + def _backend_call(self, action: str, function): # type: ignore[no-untyped-def] + self._check_thread() + try: + if self.thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return function() + return function() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to {action}" + ) from error + + def _backend_submit( + self, + op: TransferOp, + plan_handle: object, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | None, + reuse_token: object | None, + adopt_work: Callable[[BackendWork], None], + ) -> object | None: + """Call the submission SPI without allocating an extra lambda.""" + + # Both callers validate endpoint/thread state before reaching this + # same-stack helper; avoid repeating the thread-ident lookup here. + submit = cast(Callable[..., object | None], self._backend.submit) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return submit( + op, + plan_handle, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + reuse_token=reuse_token, + adopt_work=adopt_work, + ) + return submit( + op, + plan_handle, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + reuse_token=reuse_token, + adopt_work=adopt_work, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to submit {op.value}" + ) from error + + def _backend_send_notification( + self, + peer_handle: object, + payload: bytes, + *, + adopt_work: Callable[[BackendWork], None], + ) -> object | None: + """Call standalone notification send without per-call closures.""" + + # send_notification() has already validated endpoint and thread state. + send = cast(Callable[..., object | None], self._backend.send_notification) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return send(peer_handle, payload, adopt_work=adopt_work) + return send(peer_handle, payload, adopt_work=adopt_work) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to send notification" + ) from error + + def _backend_submit_prevalidated( + self, + op: TransferOp, + plan_handle: object, + *, + selection: BackendIndexSelection, + notification: bytes | None, + reuse_token: object | None, + adopt_work: Callable[[BackendWork], None], + ) -> object | None: + """Call the authenticated indexed-selection SPI without a lambda.""" + + submit = cast(Callable[..., object | None], self._backend.submit_prevalidated) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return submit( + op, + plan_handle, + selection=selection, + notification=notification, + reuse_token=reuse_token, + adopt_work=adopt_work, + ) + return submit( + op, + plan_handle, + selection=selection, + notification=notification, + reuse_token=reuse_token, + adopt_work=adopt_work, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to submit prevalidated {op.value}" + ) from error + + def _backend_submit_prevalidated_and_poll( + self, + op: TransferOp, + plan_handle: object, + *, + selection: BackendIndexSelection, + notification: bytes | None, + max_polls: int, + timeout_ns: int | None, + adopt_work: Callable[[BackendWork], None], + ) -> WorkState: + """Call the optional fused one-shot selection SPI without a lambda.""" + + submit = cast( + Callable[..., WorkState], + self._backend.submit_prevalidated_and_poll, + ) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return submit( + op, + plan_handle, + selection=selection, + notification=notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + adopt_work=adopt_work, + ) + return submit( + op, + plan_handle, + selection=selection, + notification=notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + adopt_work=adopt_work, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to submit and poll " + f"prevalidated {op.value}" + ) from error + + def _backend_create_request_slot( + self, + op: TransferOp, + plan_handle: object, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | None, + adopt_slot: Callable[..., None], + ) -> object | None: + """Call the optional persistent-request SPI without a lambda.""" + + create = cast(Callable[..., object | None], self._backend.create_request_slot) + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return create( + op, + plan_handle, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + adopt_slot=adopt_slot, + ) + return create( + op, + plan_handle, + local_indices=local_indices, + remote_indices=remote_indices, + notification=notification, + adopt_slot=adopt_slot, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to create a persistent " + f"{op.value} request slot" + ) from error + + def _backend_request_slot_start(self, start: Callable[[], None]) -> None: + """Start a cached backend slot operation on the lowest-overhead path.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = start() + else: + result = start() + if result is not None: + raise BackendFailureError("backend request-slot start must return None") + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to start a persistent request" + ) from error + + def _backend_request_slot_start_with_notification( + self, + start: Callable[[bytes | None], None], + notification: bytes | None, + ) -> None: + """Start a cached request with an explicit notification override.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = start(notification) + else: + result = start(notification) + if result is not None: + raise BackendFailureError( + "backend request-slot start_with_notification must return None" + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to start a persistent request " + "with a notification override" + ) from error + + def _backend_request_slot_start_and_poll( + self, + start_and_poll: Callable[..., WorkState | None], + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + """Enter the optional bounded persistent-request fused operation.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return start_and_poll( + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + return start_and_poll(max_polls=max_polls, timeout_ns=timeout_ns) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed during fused persistent " + "request start/poll" + ) from error + + def _backend_request_slot_start_and_poll_with_notification( + self, + start_and_poll: Callable[..., WorkState | None], + notification: bytes | None, + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + """Enter fused post/poll with a request-local notification override.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return start_and_poll( + notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + return start_and_poll( + notification, + max_polls=max_polls, + timeout_ns=timeout_ns, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed during fused persistent " + "request notification start/poll" + ) from error + + def _backend_request_slot_poll_bounded( + self, + poll_bounded: Callable[..., WorkState | None], + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState | None: + """Enter the optional bounded persistent-request polling operation.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + return poll_bounded( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed during bounded persistent " + "request polling" + ) from error + + def _backend_request_slot_state( + self, poll_state: Callable[[], WorkState | None] + ) -> WorkState | None: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return poll_state() + return poll_state() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to query persistent request state" + ) from error + + def _backend_request_slot_error( + self, backend_slot: BackendRequestSlot + ) -> BaseException | None: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return backend_slot.error + return backend_slot.error + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to query persistent request error" + ) from error + + def _backend_request_slot_completion_event( + self, backend_slot: BackendRequestSlot + ) -> object | None: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return backend_slot.completion_event + return backend_slot.completion_event + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to query persistent request " + "completion event" + ) from error + + def _backend_request_slot_cancel(self, cancel: Callable[[], bool]) -> bool: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return cancel() + return cancel() + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to cancel a persistent request" + ) from error + + def _backend_request_slot_recycle(self, recycle: Callable[[], None]) -> None: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = recycle() + else: + result = recycle() + if result is not None: + raise BackendFailureError( + "backend request-slot recycle must return None" + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to recycle a persistent request" + ) from error + + def _backend_request_slot_close(self, close: Callable[[], None]) -> None: + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = close() + else: + result = close() + if result is not None: + raise BackendFailureError("backend request-slot close must return None") + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to release a persistent request" + ) from error + + def _backend_work_state(self, backend_work: BackendWork) -> WorkState: + """Read a provider Work state without a lambda or duplicate thread check.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return backend_work.state + return backend_work.state + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to query work state" + ) from error + + def _backend_work_poll_bounded( + self, + backend_work: BackendWork, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState: + """Call a fused provider Work poll without a lambda.""" + + poll = backend_work.poll_bounded + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + return poll( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + return poll( + max_polls=max_polls, + timeout_ns=timeout_ns, + timeout_check_interval=timeout_check_interval, + ) + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to poll work" + ) from error + + def _backend_work_close(self, backend_work: BackendWork) -> None: + """Close provider Work without a bound-method/lambda dispatch wrapper.""" + + try: + if self._thread_mode is ThreadMode.SERIALIZED: + with self._backend_lock: + result = backend_work.close() + else: + result = backend_work.close() + if result is not None: + raise BackendFailureError("backend work close must return None") + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + f"backend {self.backend!r} failed to release work" + ) from error + + def _check_thread(self) -> None: + if ( + self._thread_mode is ThreadMode.SINGLE + and threading.get_ident() != self._owner_thread + ): + raise UnsupportedError( + "single-thread endpoint used from a different thread" + ) + + def __enter__(self) -> Self: + self._check_open() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +def _registered_memory( + value: object, + registration_id: str, + name: str, +) -> RegisteredMemory: + if isinstance(value, RawSpan): + return RegisteredMemory( + registration_id=registration_id, + name=name, + owner=_LifetimeAnchor(value.owner), + address=value.address, + nbytes=value.nbytes, + device=value.device, + memory_type=value.memory_type, + ) + if isinstance(value, torch.Tensor): + if value.layout is not torch.strided: + raise InvalidRegionError("only strided tensor layouts can be registered") + if not value.is_contiguous(): + raise InvalidRegionError( + "tensor registration requires a contiguous tensor; " + "no implicit copy is performed" + ) + nbytes = value.numel() * value.element_size() + address = value.data_ptr() + device = str(value.device) + memory_type = _tensor_memory_type(value) + storage = value.untyped_storage() + anchor = _LifetimeAnchor( + owner=value, + storage=storage, + allocation_address=storage.data_ptr(), + allocation_nbytes=storage.nbytes(), + ) + return _checked_memory( + registration_id, + name, + anchor, + address, + nbytes, + device, + memory_type, + ) + + untyped_storage_type = getattr(torch, "UntypedStorage", ()) + typed_storage_type = getattr(getattr(torch, "storage", object), "TypedStorage", ()) + if isinstance(value, (untyped_storage_type, typed_storage_type)): + storage = getattr(value, "_untyped_storage", value) + nbytes = storage.nbytes() + address = storage.data_ptr() + device = str(storage.device) + if storage.device.type == "cuda": + memory_type = "cuda" + else: + try: + memory_type = "host-pinned" if _storage_is_pinned(storage) else "host" + except RuntimeError: + memory_type = "host" + anchor = _LifetimeAnchor( + owner=value, + storage=storage, + allocation_address=address, + allocation_nbytes=nbytes, + ) + return _checked_memory( + registration_id, + name, + anchor, + address, + nbytes, + device, + memory_type, + ) + + try: + view = memoryview(value) # type: ignore[arg-type] + except TypeError as error: + raise InvalidRegionError( + "registration expects a contiguous Tensor, Storage, writable " + "buffer, or RawSpan" + ) from error + if not view.c_contiguous: + raise InvalidRegionError( + "buffer registration requires C-contiguous memory; " + "no implicit copy is performed" + ) + if view.readonly: + raise InvalidRegionError("buffer registration requires writable memory") + try: + address = ctypes.addressof(ctypes.c_char.from_buffer(view)) + except (TypeError, ValueError) as error: + raise InvalidRegionError("could not obtain the buffer address") from error + anchor = _LifetimeAnchor( + owner=value, + buffer=view, + allocation_address=address, + allocation_nbytes=view.nbytes, + ) + return _checked_memory( + registration_id, + name, + anchor, + address, + view.nbytes, + "cpu", + "host", + ) + + +def _checked_memory( + registration_id: str, + name: str, + owner: object, + address: int, + nbytes: int, + device: str, + memory_type: str, +) -> RegisteredMemory: + if not isinstance(address, int) or address <= 0 or nbytes <= 0: + raise InvalidRegionError("cannot register empty or unallocated memory") + return RegisteredMemory( + registration_id=registration_id, + name=name, + owner=owner, + address=address, + nbytes=nbytes, + device=device, + memory_type=memory_type, + ) + + +def _released_memory(memory: RegisteredMemory) -> RegisteredMemory: + return RegisteredMemory( + registration_id=memory.registration_id, + name=memory.name, + owner=None, + address=memory.address, + nbytes=memory.nbytes, + device=memory.device, + memory_type=memory.memory_type, + ) + + +def _memory_ranges_overlap(first: RegisteredMemory, second: RegisteredMemory) -> bool: + if first.device != second.device: + return False + return ( + first.address < second.address + second.nbytes + and second.address < first.address + first.nbytes + ) + + +def _validate_memory_stable(memory: RegisteredMemory) -> None: + anchor = memory.owner + if not isinstance(anchor, _LifetimeAnchor): + return + if anchor.storage is not None: + storage = cast(Any, anchor.storage) + try: + address = storage.data_ptr() + nbytes = storage.nbytes() + except Exception as error: + raise InvalidRegionError("could not validate registered Storage") from error + if address != anchor.allocation_address or nbytes != anchor.allocation_nbytes: + raise InvalidRegionError( + "registered Storage was resized; deregister and register it again" + ) + if anchor.buffer is not None: + try: + address = ctypes.addressof(ctypes.c_char.from_buffer(anchor.buffer)) + except (TypeError, ValueError) as error: + raise InvalidRegionError( + "registered buffer no longer exposes its original address" + ) from error + if ( + address != anchor.allocation_address + or anchor.buffer.nbytes != anchor.allocation_nbytes + ): + raise InvalidRegionError("registered buffer address or size changed") + + +def _tensor_memory_type(tensor: torch.Tensor) -> str: + if tensor.device.type == "cuda": + return "cuda" + if tensor.device.type == "cpu": + try: + return "host-pinned" if tensor.is_pinned() else "host" + except RuntimeError: + return "host" + return tensor.device.type + + +def _storage_is_pinned(storage: object) -> bool: + view = torch.empty( + 0, + dtype=torch.uint8, + device=storage.device, # type: ignore[attr-defined] + ) + return view.set_(storage).is_pinned() # type: ignore[arg-type] + + +def _validate_region_bounds( + allocation_nbytes: int, + offset: int, + nbytes: int | None, + stride: int | None, + count: int | None, +) -> tuple[int, int, int | None, int | None]: + _validate_nonnegative_int(offset, "offset") + stride, count = _validate_stride(stride, count) + if nbytes is None: + if stride is not None: + raise InvalidRegionError( + "nbytes is required when stride and count are specified" + ) + nbytes = allocation_nbytes - offset + _validate_positive_int(nbytes, "nbytes") + if stride is not None and stride < nbytes: + raise InvalidRegionError("stride must be at least nbytes") + extent = _region_extent(nbytes, stride, count) + if offset + extent > allocation_nbytes: + raise InvalidRegionError( + f"region [{offset}, {offset + extent}) exceeds allocation size " + f"{allocation_nbytes}" + ) + return offset, nbytes, stride, count + + +def _validate_stride( + stride: int | None, + count: int | None, +) -> tuple[int | None, int | None]: + if (stride is None) != (count is None): + raise InvalidRegionError("stride and count must be specified together") + if stride is not None: + _validate_positive_int(stride, "stride") + _validate_positive_int(count, "count") + return stride, count + + +def _region_extent( + nbytes: int, + stride: int | None, + count: int | None, +) -> int: + if stride is None: + return nbytes + if count is None: + raise InvalidRegionError("count is required with stride") + return (count - 1) * stride + nbytes + + +def _validate_plan_indices( + plan: TransferPlan, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + op: TransferOp, +) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]: + if not plan.indexed: + if local_indices is not None or remote_indices is not None: + raise InvalidRegionError( + "indices require a plan prepared with indexed=True" + ) + local = remote = None + else: + if (local_indices is None) != (remote_indices is None): + raise InvalidRegionError( + "local_indices and remote_indices must be specified together" + ) + if local_indices is None: + local = remote = None + else: + remote_values = cast(Sequence[int], remote_indices) + local = _validate_indices( + local_indices, plan.local_block_count, "local_indices" + ) + remote = _validate_indices( + remote_values, + plan.remote_block_count, + "remote_indices", + ) + if local is None: + if plan._local_total_nbytes != plan._remote_total_nbytes: + raise InvalidRegionError( + f"selected local and remote byte totals differ: " + f"{plan._local_total_nbytes} != {plan._remote_total_nbytes}" + ) + destination_nonoverlap = ( + plan._remote_full_nonoverlap + if op is TransferOp.WRITE + else plan._local_full_nonoverlap + ) + if not destination_nonoverlap: + raise InvalidRegionError("selected destination regions must not overlap") + return local, remote + + selected_local = ( + None + if local is None + else _select_region_blocks(plan._local, local, plan._local_starts) + ) + selected_remote = ( + None + if remote is None + else _select_region_blocks(plan._remote, remote, plan._remote_starts) + ) + local_size = _selected_region_nbytes(plan._local, selected_local) + remote_size = _selected_region_nbytes(plan._remote, selected_remote) + if local_size != remote_size: + raise InvalidRegionError( + f"selected local and remote byte totals differ: " + f"{local_size} != {remote_size}" + ) + if op is TransferOp.WRITE: + _validate_nonoverlapping_destination(plan._remote, selected_remote) + else: + _validate_nonoverlapping_destination(plan._local, selected_local) + return local, remote + + +def _snapshot_backend_index_selection( + plan: TransferPlan, + op: TransferOp, + local_indices: Sequence[int], + remote_indices: Sequence[int], + *, + collect_local_devices: bool = True, +) -> tuple[BackendIndexSelection, tuple[str, ...]]: + """Create the sole private immutable selection from one input snapshot.""" + + ( + packed, + local_count, + remote_count, + selected_nbytes, + uniform_pair_layout, + destination_catalog_nonoverlapping, + local_devices, + ) = _pack_and_validate_index_selection( + plan, + op, + local_indices, + remote_indices, + collect_local_devices=collect_local_devices, + ) + # ``packed`` is owned immutable bytes, so the cast is already read-only. + # Avoid allocating a second memoryview wrapper on every dynamic selection. + packed_view = memoryview(packed).cast("i") + local_view = packed_view[:local_count] + remote_view = packed_view[local_count : local_count + remote_count] + return ( + _make_backend_index_selection( + op, + plan._handle, + local_view, + remote_view, + selected_nbytes=selected_nbytes, + uniform_pair_layout=uniform_pair_layout, + destination_catalog_nonoverlapping=destination_catalog_nonoverlapping, + destination_indices_unique=True, + ), + local_devices, + ) + + +def _pack_and_validate_index_selection( + plan: TransferPlan, + op: TransferOp, + local_indices: Sequence[int], + remote_indices: Sequence[int], + *, + collect_local_devices: bool = True, +) -> tuple[bytes, int, int, int, bool, bool, tuple[str, ...]]: + """Snapshot two vectors once and derive reusable Core proof facts.""" + + local_input_view = _native_int32_buffer_view(local_indices) + remote_input_view = _native_int32_buffer_view(remote_indices) + native_destination_duplicate: bool | None = None + if local_input_view is not None and remote_input_view is not None: + local_count = local_input_view.nbytes // 4 + if local_count == 0: + raise InvalidRegionError("local_indices must be non-empty") + remote_count = remote_input_view.nbytes // 4 + if remote_count == 0: + raise InvalidRegionError("remote_indices must be non-empty") + + if _compiled_index_analyzer is not None: + # The helper owns the only copy: it creates one combined PyBytes + # snapshot under the GIL, then derives every fact from that + # immutable snapshot with the GIL released. + ( + packed, + analyzed_local_count, + analyzed_remote_count, + local_minimum, + local_maximum, + remote_minimum, + remote_maximum, + native_destination_duplicate, + ) = _compiled_index_analyzer( + local_input_view, + remote_input_view, + op is TransferOp.WRITE, + ) + if ( + analyzed_local_count != local_count + or analyzed_remote_count != remote_count + ): + raise RuntimeError( + "compiled int32 index analyzer returned inconsistent counts" + ) + local_bounds = (local_minimum, local_maximum) + remote_bounds = (remote_minimum, remote_maximum) + else: + # Compatibility path for source-tree Python loaded with an older + # torch._C binary. It preserves the exact pre-helper behavior and + # is deliberately selected only by a missing import, never by a + # helper exception. + packed = b"".join((local_input_view.cast("B"), remote_input_view.cast("B"))) + packed_indices = memoryview(packed).cast("i") + local_view = packed_indices[:local_count] + remote_view = packed_indices[local_count : local_count + remote_count] + destination_view = remote_view if op is TransferOp.WRITE else local_view + if max(local_count, remote_count) <= _INDEX_SELECTION_TORCH_THRESHOLD: + local_bounds = (min(local_view), max(local_view)) + remote_bounds = (min(remote_view), max(remote_view)) + native_destination_duplicate = len(set(destination_view)) != len( + destination_view + ) + else: + validation_scratch = bytearray(packed) + validation = torch.frombuffer(validation_scratch, dtype=torch.int32) + local_tensor = validation[:local_count] + remote_tensor = validation[local_count : local_count + remote_count] + local_minimum, local_maximum = torch.aminmax(local_tensor) + remote_minimum, remote_maximum = torch.aminmax(remote_tensor) + destination_tensor = ( + remote_tensor if op is TransferOp.WRITE else local_tensor + ) + ordered_destination = torch.sort(destination_tensor).values + local_bounds = (int(local_minimum), int(local_maximum)) + remote_bounds = (int(remote_minimum), int(remote_maximum)) + native_destination_duplicate = bool( + torch.any(ordered_destination[1:] == ordered_destination[:-1]) + ) + + packed_indices = memoryview(packed).cast("i") + local_view = packed_indices[:local_count] + remote_view = packed_indices[local_count : local_count + remote_count] + + # Preserve the public error order: local bounds, remote bounds, byte + # equality, then destination duplicate/physical-overlap diagnostics. + _raise_packed_index_bounds( + *local_bounds, plan.local_block_count, "local_indices" + ) + _raise_packed_index_bounds( + *remote_bounds, plan.remote_block_count, "remote_indices" + ) + else: + # Mixed buffers and general Sequences retain the established + # operator.index normalization and validation behavior. + packed = bytearray() + local_count, local_needs_bounds = _append_index_snapshot( + packed, + local_indices, + plan.local_block_count, + "local_indices", + ) + if local_needs_bounds: + # Preserve the public left-to-right validation contract for mixed + # inputs. A native local buffer is copied by _append_index_snapshot, + # but its bounds must be diagnosed before the remote Sequence is + # iterated (which may itself be invalid or have side effects). + local_snapshot_view = memoryview(packed).cast("i") + try: + _validate_packed_index_bounds( + local_snapshot_view, + plan.local_block_count, + "local_indices", + ) + finally: + # bytearray cannot be resized while an exported view exists. + local_snapshot_view.release() + remote_count, remote_needs_bounds = _append_index_snapshot( + packed, + remote_indices, + plan.remote_block_count, + "remote_indices", + ) + packed_indices = memoryview(packed).cast("i") + local_view = packed_indices[:local_count] + remote_view = packed_indices[local_count : local_count + remote_count] + if remote_needs_bounds: + _validate_packed_index_bounds( + remote_view, plan.remote_block_count, "remote_indices" + ) + + selected_local: tuple[tuple[Region, int], ...] | None = None + selected_remote: tuple[tuple[Region, int], ...] | None = None + if plan._local_uniform_nbytes is None: + selected_local = _select_region_blocks( + plan._local, local_view, plan._local_starts + ) + local_size = _selected_region_nbytes(plan._local, selected_local) + else: + local_size = local_count * plan._local_uniform_nbytes + if plan._remote_uniform_nbytes is None: + selected_remote = _select_region_blocks( + plan._remote, remote_view, plan._remote_starts + ) + remote_size = _selected_region_nbytes(plan._remote, selected_remote) + else: + remote_size = remote_count * plan._remote_uniform_nbytes + if local_size != remote_size: + raise InvalidRegionError( + f"selected local and remote byte totals differ: " + f"{local_size} != {remote_size}" + ) + + destination_full_nonoverlap = ( + plan._remote_full_nonoverlap + if op is TransferOp.WRITE + else plan._local_full_nonoverlap + ) + if native_destination_duplicate: + raise InvalidRegionError("selected destination regions must not overlap") + if destination_full_nonoverlap: + destination_view = remote_view if op is TransferOp.WRITE else local_view + if native_destination_duplicate is None: + _validate_unique_packed_indices(destination_view) + elif op is TransferOp.WRITE: + if selected_remote is None: + selected_remote = _select_region_blocks( + plan._remote, remote_view, plan._remote_starts + ) + _validate_nonoverlapping_destination(plan._remote, selected_remote) + else: + if selected_local is None: + selected_local = _select_region_blocks( + plan._local, local_view, plan._local_starts + ) + _validate_nonoverlapping_destination(plan._local, selected_local) + + if not collect_local_devices: + local_devices = () + elif plan._local_uniform_device is None: + if selected_local is None: + selected_local = _select_region_blocks( + plan._local, local_view, plan._local_starts + ) + local_devices = tuple( + dict.fromkeys(region.device for region, _ in selected_local) + ) + else: + local_devices = (plan._local_uniform_device,) + + uniform_pair_layout = ( + local_count == remote_count + and plan._local_uniform_nbytes is not None + and plan._local_uniform_nbytes == plan._remote_uniform_nbytes + ) + return ( + bytes(packed), + local_count, + remote_count, + local_size, + uniform_pair_layout, + destination_full_nonoverlap, + local_devices, + ) + + +def _append_index_snapshot( + packed: bytearray, + values: Sequence[int], + upper_bound: int, + name: str, +) -> tuple[int, bool]: + """Append one vector, combining general-sequence coercion and bounds.""" + + if values is None: + raise InvalidRegionError(f"{name} must be non-empty") + view = _native_int32_buffer_view(values) + if view is not None: + count = view.nbytes // 4 + if count == 0: + raise InvalidRegionError(f"{name} must be non-empty") + packed.extend(view.cast("B")) + return count, True + + count = 0 + try: + iterator = iter(values) + except TypeError as error: + raise InvalidRegionError(f"{name} must contain integers") from error + for raw_value in iterator: + if isinstance(raw_value, bool): + raise InvalidRegionError(f"{name} must contain integers") + try: + value = operator.index(raw_value) + except TypeError as error: + raise InvalidRegionError(f"{name} must contain integers") from error + if value < 0 or value >= upper_bound: + raise InvalidRegionError( + f"{name} index {value} is outside [0, {upper_bound})" + ) + if value > _INT32_MAX: + raise InvalidRegionError( + f"{name} index {value} exceeds the packed int32 range" + ) + packed.extend(value.to_bytes(4, byteorder=sys.byteorder, signed=True)) + count += 1 + if count == 0: + raise InvalidRegionError(f"{name} must be non-empty") + return count, False + + +def _native_int32_buffer_view(values: object) -> memoryview | None: + """Return the exact packed-input view, without copying or coercion.""" + + try: + view = memoryview(values) + except TypeError: + return None + if ( + view.ndim == 1 + and view.c_contiguous + and view.itemsize == 4 + and view.format in _NATIVE_INT32_FORMATS + ): + return view + return None + + +def _raise_packed_index_bounds( + low: int, + high: int, + upper_bound: int, + name: str, +) -> None: + if low < 0: + raise InvalidRegionError(f"{name} index {low} is outside [0, {upper_bound})") + if high >= upper_bound: + raise InvalidRegionError(f"{name} index {high} is outside [0, {upper_bound})") + + +def _validate_packed_index_bounds( + values: memoryview, + upper_bound: int, + name: str, +) -> None: + tensor = torch.frombuffer(values, dtype=torch.int32) + minimum, maximum = torch.aminmax(tensor) + low = int(minimum) + high = int(maximum) + _raise_packed_index_bounds(low, high, upper_bound, name) + + +def _validate_unique_packed_indices(values: memoryview) -> None: + if len(values) < 2: + return + tensor = torch.frombuffer(values, dtype=torch.int32) + ordered = torch.sort(tensor).values + if bool(torch.any(ordered[1:] == ordered[:-1])): + raise InvalidRegionError("selected destination regions must not overlap") + + +def _validate_cuda_ordering_arguments( + capabilities: Capabilities, + backend: str, + mode: CudaOrderingMode, + stream: object | None, + ready_event: object | None, +) -> None: + if stream is not None and ready_event is not None: + raise ValueError("stream and ready_event are mutually exclusive") + ordering_token_count = int(stream is not None) + int(ready_event is not None) + if mode is CudaOrderingMode.CALLER_READY: + if ordering_token_count: + raise ValueError( + "caller_ready CUDA ordering does not accept a stream or ready_event" + ) + elif ordering_token_count != 1: + raise ValueError( + "host_wait CUDA ordering requires exactly one stream or ready_event" + ) + if mode not in capabilities.cuda_ordering_modes: + raise UnsupportedError( + f"backend {backend!r} does not support CUDA ordering " + f"mode {mode.value!r}" + ) + + +def _validate_attached_notification( + capabilities: Capabilities, + backend: str, + notification: bytes | bytearray | memoryview | None, +) -> bytes | None: + if notification is None: + return None + if not capabilities.attached_notifications: + raise UnsupportedError(f"backend {backend!r} does not support notifications") + result = _as_bytes( + notification, + "notification", + max_nbytes=_MAX_NOTIFICATION_BYTES, + ) + if not result: + raise ValueError("an attached notification payload must be non-empty") + return result + + +def _validate_request_slot_notification( + notification: bytes | bytearray | memoryview | None, +) -> bytes | None: + """Validate an explicit per-generation slot notification or clear.""" + + if notification is None: + return None + result = _as_bytes( + notification, + "notification", + max_nbytes=_MAX_NOTIFICATION_BYTES, + ) + if not result: + raise ValueError("an attached notification payload must be non-empty") + return result + + +def _validate_cuda_ordering_selection( + plan: TransferPlan, + local_indices: Sequence[int] | None, + mode: CudaOrderingMode, + stream: object | None, + ready_event: object | None, +) -> torch.device | None: + """Validate a CUDA submission shape without establishing readiness.""" + + if mode is CudaOrderingMode.CALLER_READY: + return None + if mode is not CudaOrderingMode.HOST_WAIT: + raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") + + selected = ( + tuple((region, 0) for region in plan._local) + if local_indices is None + else _select_region_blocks(plan._local, local_indices, plan._local_starts) + ) + return _validate_cuda_ordering_devices( + tuple(dict.fromkeys(region.device for region, _ in selected)), + mode, + stream, + ready_event, + ) + + +def _validate_cuda_ordering_devices( + selected_devices: Sequence[str], + mode: CudaOrderingMode, + stream: object | None, + ready_event: object | None, +) -> torch.device | None: + """Validate ordering from device facts cached by an IndexSelection.""" + + if mode is CudaOrderingMode.CALLER_READY: + return None + if mode is not CudaOrderingMode.HOST_WAIT: + raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") + + devices: set[torch.device] = set() + for device_name in selected_devices: + try: + device = torch.device(device_name) + except (RuntimeError, TypeError, ValueError) as error: + raise InvalidRegionError( + f"local region has invalid device {device_name!r}" + ) from error + if device.type != "cuda" or device.index is None: + raise InvalidRegionError( + "host_wait requires selected local regions on one explicitly " + "indexed CUDA device" + ) + devices.add(device) + if len(devices) != 1: + raise UnsupportedError( + "host_wait requires selected local regions on exactly one CUDA device" + ) + expected_device = next(iter(devices)) + + if stream is not None: + stream_types = tuple( + stream_type + for stream_type in ( + getattr(torch, "Stream", None), + getattr(torch.cuda, "Stream", None), + ) + if isinstance(stream_type, type) + ) + if not stream_types or not isinstance(stream, stream_types): + raise TypeError("stream must be a torch CUDA Stream") + _validate_cuda_token_device(stream, expected_device, "stream") + else: + event_type = getattr(torch.cuda, "Event", None) + if not isinstance(event_type, type) or not isinstance(ready_event, event_type): + raise TypeError("ready_event must be a torch.cuda.Event") + _validate_cuda_token_device(ready_event, expected_device, "ready_event") + + return expected_device + + +def _establish_cuda_ordering_prevalidated( + mode: CudaOrderingMode, + stream: object | None, + ready_event: object | None, + expected_device: torch.device | None, +) -> None: + """Resolve Core-owned CUDA readiness before entering provider locks.""" + + if mode is CudaOrderingMode.CALLER_READY: + return + if mode is not CudaOrderingMode.HOST_WAIT or expected_device is None: + raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") + + if stream is not None: + try: + if _cuda_token_stream_is_capturing(stream): + raise UnsupportedError("host_wait is not CUDA Graph capture safe") + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to query CUDA Graph capture state" + ) from error + + if stream is not None: + try: + event = stream.record_event() # type: ignore[attr-defined] + except Exception as error: + raise BackendFailureError( + "failed to record the CUDA readiness event" + ) from error + else: + event = ready_event + + if stream is not None: + # A stream implementation supplies the event on each submission; keep + # this runtime check even though the stream itself was validated once. + event_type = getattr(torch.cuda, "Event", None) + if not isinstance(event_type, type) or not isinstance(event, event_type): + raise TypeError("ready_event must be a torch.cuda.Event") + _validate_cuda_token_device(event, expected_device, "ready_event") + if event is None: + raise BackendFailureError("CUDA readiness event is missing") + try: + event.synchronize() # type: ignore[union-attr] + except TransferError: + raise + except Exception as error: + raise BackendFailureError("failed to wait for CUDA readiness") from error + + +def _cuda_token_stream_is_capturing(stream: object) -> bool: + """Query the supplied stream before Core records a new readiness event.""" + + direct_query = getattr(stream, "is_capturing", None) + if callable(direct_query): + return bool(direct_query()) + with torch.cuda.stream(stream): # type: ignore[arg-type] + return torch.cuda.is_current_stream_capturing() + + +def _make_reusable_cuda_stream_readiness( + stream: object, + expected_device: torch.device | None, +) -> Callable[[], None]: + """Create one per-slot event and return its generation readiness gate.""" + + if expected_device is None: + raise UnsupportedError("host_wait stream requires one CUDA device") + try: + event = torch.cuda.Event() + record = cast(Callable[[object], None], event.record) + synchronize = cast(Callable[[], None], event.synchronize) + except Exception as error: + raise BackendFailureError( + "failed to create the reusable CUDA readiness event" + ) from error + + direct_capture_query = getattr(stream, "is_capturing", None) + if callable(direct_capture_query): + capture_query = direct_capture_query + else: + stream_context = torch.cuda.stream + current_stream_is_capturing = torch.cuda.is_current_stream_capturing + + def capture_query() -> bool: + with stream_context(stream): # type: ignore[arg-type] + return bool(current_stream_is_capturing()) + + device_validated = False + + def synchronize_generation() -> None: + nonlocal device_validated + try: + if capture_query(): + raise UnsupportedError("host_wait is not CUDA Graph capture safe") + except TransferError: + raise + except Exception as error: + raise BackendFailureError( + "failed to query CUDA Graph capture state" + ) from error + try: + record(stream) + except Exception as error: + raise BackendFailureError( + "failed to record the CUDA readiness event" + ) from error + if not device_validated: + _validate_cuda_token_device(event, expected_device, "ready_event") + device_validated = True + try: + synchronize() + except TransferError: + raise + except Exception as error: + raise BackendFailureError("failed to wait for CUDA readiness") from error + + return synchronize_generation + + +def _validate_cuda_token_device( + token: object, + expected: torch.device, + name: str, +) -> None: + device_value = getattr(token, "device", None) + if device_value is None: + qualifier = "recorded " if name == "ready_event" else "" + raise InvalidRegionError(f"{name} must be a {qualifier}CUDA token") + try: + device = torch.device(device_value) + except (RuntimeError, TypeError, ValueError) as error: + raise InvalidRegionError(f"{name} has an invalid device") from error + if device != expected: + raise InvalidRegionError( + f"{name} device {device} does not match local region device {expected}" + ) + + +def _logical_block_count(regions: Sequence[Region]) -> int: + return sum(region.count or 1 for region in regions) + + +def _uniform_region_nbytes(regions: Sequence[Region]) -> int | None: + first = regions[0].nbytes + return first if all(region.nbytes == first for region in regions[1:]) else None + + +def _uniform_region_device(regions: Sequence[Region]) -> str | None: + first = regions[0].device + return first if all(region.device == first for region in regions[1:]) else None + + +def _catalog_starts(regions: Sequence[Region]) -> tuple[int, ...]: + starts: list[int] = [] + next_index = 0 + for region in regions: + starts.append(next_index) + next_index += region.count or 1 + return tuple(starts) + + +def _select_region_blocks( + regions: Sequence[Region], + indices: Sequence[int], + starts: Sequence[int], +) -> tuple[tuple[Region, int], ...]: + selected: list[tuple[Region, int]] = [] + for index in indices: + region_index = bisect_right(starts, index) - 1 + if region_index < 0 or region_index >= len(regions): + raise InvalidRegionError( + "logical block index is outside the region catalog" + ) + selected.append((regions[region_index], index - starts[region_index])) + return tuple(selected) + + +def _selected_region_nbytes( + regions: Sequence[Region], + selected: Sequence[tuple[Region, int]] | None, +) -> int: + if selected is None: + return sum(region.total_nbytes for region in regions) + return sum(region.nbytes for region, _ in selected) + + +def _validate_nonoverlapping_destination( + regions: Sequence[Region], + selected: Sequence[tuple[Region, int]] | None, +) -> None: + intervals: dict[object, list[tuple[int, int]]] = {} + blocks = ( + ( + (region, block_index) + for region in regions + for block_index in range(region.count or 1) + ) + if selected is None + else iter(selected) + ) + for region, block_index in blocks: + descriptor = region._backend_descriptor + if isinstance(region._origin, Registration): + key: object = ("local", region.device) + base = region._origin._memory.address + elif isinstance(descriptor, RegistrationMetadata): + key = ( + "remote-registration", + id(region._origin), + descriptor.registration_id, + ) + base = 0 + elif isinstance(descriptor, _RemoteRawDescriptor): + key = ( + "remote-address", + id(region._origin), + descriptor.device, + descriptor.memory_type, + ) + base = descriptor.address + else: + raise InvalidRegionError("region has an invalid memory descriptor") + stride = region.stride or region.nbytes + start = base + region.offset + block_index * stride + intervals.setdefault(key, []).append((start, start + region.nbytes)) + for values in intervals.values(): + values.sort() + for previous, current in itertools.pairwise(values): + if current[0] < previous[1]: + raise InvalidRegionError( + "selected destination regions must not overlap" + ) + + +def _regions_are_nonoverlapping(regions: Sequence[Region]) -> bool: + # Each compact Region is internally disjoint because stride >= nbytes. + # Compare only its bounding envelope against other regions in the same + # physical address space. An ambiguous interleaving is conservatively + # rejected instead of expanding `count` logical blocks. + envelopes: dict[object, list[tuple[int, int]]] = {} + for region in regions: + descriptor = region._backend_descriptor + if isinstance(region._origin, Registration): + key: object = ("local", region.device) + base = region._origin._memory.address + elif isinstance(descriptor, RegistrationMetadata): + key = ( + "remote-registration", + id(region._origin), + descriptor.registration_id, + ) + base = 0 + elif isinstance(descriptor, _RemoteRawDescriptor): + key = ( + "remote-address", + id(region._origin), + descriptor.device, + descriptor.memory_type, + ) + base = descriptor.address + else: + return False + stride = region.stride or region.nbytes + start = base + region.offset + end = start + ((region.count or 1) - 1) * stride + region.nbytes + envelopes.setdefault(key, []).append((start, end)) + for values in envelopes.values(): + values.sort() + for previous, current in itertools.pairwise(values): + if current[0] < previous[1]: + return False + return True + + +def _validate_indices( + values: Sequence[int], + upper_bound: int, + name: str, +) -> tuple[int, ...]: + result: list[int] = [] + for raw_value in values: + if isinstance(raw_value, bool): + raise InvalidRegionError(f"{name} must contain integers") + try: + value = operator.index(raw_value) + except TypeError as error: + raise InvalidRegionError(f"{name} must contain integers") from error + if value < 0 or value >= upper_bound: + raise InvalidRegionError( + f"{name} index {value} is outside [0, {upper_bound})" + ) + result.append(value) + if not result: + raise InvalidRegionError(f"{name} must be non-empty") + return tuple(result) + + +def _adopt_resource_handle(resource: object, handle: object, kind: str) -> None: + if handle is None or handle is _UNADOPTED_HANDLE: + raise BackendFailureError(f"backend adopted an invalid {kind} handle") + mutable_resource = cast(Any, resource) + current = mutable_resource._handle + if current is _UNADOPTED_HANDLE: + mutable_resource._handle = handle + return + if current is handle: + return + raise BackendFailureError(f"backend attempted to adopt {kind} more than once") + + +def _creation_cleanup_failure( + kind: str, + original_error: BaseException, + cleanup_error: BaseException, +) -> BackendFailureError: + del cleanup_error + failure = BackendFailureError( + f"{kind} creation failed and cleanup is pending; recover it from the " + f"Endpoint open-resource snapshots and retry close()" + ) + add_note = getattr(failure, "add_note", None) + if callable(add_note): + add_note(f"original failure type: {type(original_error).__name__}") + return failure + + +def _endpoint_creation_cleanup_failure( + endpoint: Endpoint, + original_error: BaseException, + cleanup_error: BaseException, +) -> BackendFailureError: + failure = BackendFailureError( + "endpoint backend initialization failed and cleanup is pending; recover " + "exception.recovery_endpoint and retry close()" + ) + failure.recovery_endpoint = endpoint # type: ignore[attr-defined] + add_note = getattr(failure, "add_note", None) + if callable(add_note): + add_note(f"initialization failure type: {type(original_error).__name__}") + add_note(f"cleanup failure type: {type(cleanup_error).__name__}") + return failure + + +def _as_bytes( + value: object, + name: str, + *, + max_nbytes: int | None = None, +) -> bytes: + if type(value) is bytes: + if max_nbytes is not None and len(value) > max_nbytes: + raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") + return value + if not isinstance(value, (bytes, bytearray, memoryview)): + raise TypeError(f"{name} must be bytes-like") + if max_nbytes is not None: + try: + source_nbytes = ( + value.nbytes if isinstance(value, memoryview) else len(value) + ) + except ValueError as error: + raise ValueError(f"{name} buffer is unavailable") from error + if source_nbytes > max_nbytes: + raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") + result = bytes(value) + if max_nbytes is not None and len(result) > max_nbytes: + raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") + return result + + +def _validate_positive_int(value: object, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise InvalidRegionError(f"{name} must be a positive integer") + + +def _transfer_slot_deadline_ns(timeout: float | None) -> int | None: + if timeout is None: + return None + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): + raise TypeError("timeout must be a non-negative number or None") + if timeout < 0 or (isinstance(timeout, float) and not math.isfinite(timeout)): + raise ValueError("timeout must be finite and non-negative") + if isinstance(timeout, int): + timeout_ns = timeout * 1_000_000_000 + else: + numerator, denominator = timeout.as_integer_ratio() + timeout_ns = (numerator * 1_000_000_000 + denominator - 1) // denominator + return time.monotonic_ns() + timeout_ns + + +def _validate_bounded_poll_arguments( + *, + max_polls: object, + timeout_ns: object, + timeout_check_interval: object | None, + allow_zero: bool, +) -> None: + minimum = 0 if allow_zero else 1 + if isinstance(max_polls, bool) or not isinstance(max_polls, int): + qualifier = "non-negative" if allow_zero else "positive" + raise TypeError(f"max_polls must be a {qualifier} integer") + if max_polls < minimum: + qualifier = "non-negative" if allow_zero else "positive" + raise ValueError(f"max_polls must be a {qualifier} integer") + if max_polls > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS: + qualifier = "non-negative" if allow_zero else "positive" + raise ValueError(f"max_polls must be a {qualifier} signed 64-bit integer") + if timeout_ns is not None and ( + isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int) + ): + raise TypeError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + if timeout_ns is not None and ( + timeout_ns < 0 or timeout_ns > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS + ): + raise ValueError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + if timeout_check_interval is not None and ( + isinstance(timeout_check_interval, bool) + or not isinstance(timeout_check_interval, int) + ): + raise TypeError("timeout_check_interval must be a positive integer") + if timeout_check_interval is not None and timeout_check_interval <= 0: + raise ValueError("timeout_check_interval must be a positive integer") + if ( + timeout_check_interval is not None + and timeout_check_interval > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS + ): + raise ValueError( + "timeout_check_interval must be a positive signed 64-bit integer" + ) + + +def _validate_nonnegative_int(value: object, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise InvalidRegionError(f"{name} must be a non-negative integer") + + +def _validate_nonempty_str( + value: object, + name: str, + error_type: type[Exception] = ValueError, +) -> None: + if not isinstance(value, str) or not value: + raise error_type(f"{name} must be a non-empty string") + + +def _coerce_enum(value: object, enum_type: type[Any], name: str) -> Any: + if isinstance(value, enum_type): + return value + try: + return enum_type(value) + except (TypeError, ValueError) as error: + choices = ", ".join(repr(item.value) for item in enum_type) + raise ValueError(f"{name} must be one of {choices}") from error diff --git a/torch/distributed/_transfer/_backend.py b/torch/distributed/_transfer/_backend.py new file mode 100644 index 0000000000000..49a77b0e737e1 --- /dev/null +++ b/torch/distributed/_transfer/_backend.py @@ -0,0 +1,1034 @@ +from __future__ import annotations + +import threading +import time +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol + +from ._errors import BackendNotFoundError, BusyError, UnsupportedError +from ._types import ( + BackendNotification, + BackendNotificationBatch, + BackendRegion, + Capabilities, + EndpointMetadata, + ProgressMode, + ThreadMode, + TransferOp, + WorkState, +) + + +__all__ = [ + "BACKEND_FACTORY_API_VERSION", + "Backend", + "BackendFactory", + "BackendFactoryV2", + "BackendIndexSelection", + "BackendRequestSlot", + "BackendWork", + "ImportedRawMemory", + "ImportedRegistration", + "RegisteredMemory", + "get_backend_factory", + "register_backend", + "unregister_backend", +] + +_REQUEST_SLOT_TERMINAL_STATES = frozenset( + {WorkState.COMPLETED, WorkState.FAILED, WorkState.CANCELLED} +) +_SIGNED_INT64_MAX = (1 << 63) - 1 +_INDEX_SELECTION_TOKEN = object() + + +class BackendIndexSelection: + """Core-authenticated, immutable indexed-selection lowering input. + + Instances are constructed only by Core after it has snapshotted and + validated both index vectors against one exact plan and operation. The + read-only memoryviews share one private, Core-owned native-int32 buffer. + Providers must validate the exact type and plan/operation binding before + using the proof flags; :meth:`Backend.submit_prevalidated` does that for + source-compatible providers. + """ + + __slots__ = ( + "_destination_catalog_nonoverlapping", + "_destination_indices_unique", + "_local_indices", + "_op", + "_plan_handle", + "_remote_indices", + "_sealed", + "_selected_nbytes", + "_uniform_pair_layout", + ) + + def __init__( + self, + op: TransferOp, + plan_handle: object, + local_indices: memoryview, + remote_indices: memoryview, + *, + selected_nbytes: int, + uniform_pair_layout: bool, + destination_catalog_nonoverlapping: bool, + destination_indices_unique: bool, + _token: object, + ) -> None: + if _token is not _INDEX_SELECTION_TOKEN: + raise TypeError("BackendIndexSelection objects are created by Core") + # This object is created on every full-dynamic submission. Calling its + # rejecting Python __setattr__ for each field would repeat getattr and + # dispatch overhead nine times before the object is sealed. + set_attribute = object.__setattr__ + set_attribute(self, "_op", op) + set_attribute(self, "_plan_handle", plan_handle) + set_attribute(self, "_local_indices", local_indices) + set_attribute(self, "_remote_indices", remote_indices) + set_attribute(self, "_selected_nbytes", selected_nbytes) + set_attribute(self, "_uniform_pair_layout", uniform_pair_layout) + set_attribute( + self, + "_destination_catalog_nonoverlapping", + destination_catalog_nonoverlapping, + ) + set_attribute(self, "_destination_indices_unique", destination_indices_unique) + set_attribute(self, "_sealed", True) + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError("BackendIndexSelection is immutable") + + def __delattr__(self, name: str) -> None: + raise AttributeError("BackendIndexSelection is immutable") + + @property + def local_indices(self) -> memoryview: + return self._local_indices + + @property + def remote_indices(self) -> memoryview: + return self._remote_indices + + @property + def selected_nbytes(self) -> int: + return self._selected_nbytes + + @property + def uniform_pair_layout(self) -> bool: + """Whether Core proved equal, uniform per-pair block geometry.""" + + return self._uniform_pair_layout + + @property + def destination_catalog_nonoverlapping(self) -> bool: + """Whether Core proved that the full destination catalog is disjoint.""" + + return self._destination_catalog_nonoverlapping + + @property + def destination_indices_unique(self) -> bool: + """Whether Core proved the destination logical indices are unique.""" + + return self._destination_indices_unique + + def _validate_binding(self, op: TransferOp, plan_handle: object) -> None: + if type(self) is not BackendIndexSelection: + raise TypeError("selection must be a Core BackendIndexSelection") + if self._op is not op or self._plan_handle is not plan_handle: + raise ValueError("prevalidated selection does not match plan and op") + + +def _make_backend_index_selection( + op: TransferOp, + plan_handle: object, + local_indices: memoryview, + remote_indices: memoryview, + *, + selected_nbytes: int, + uniform_pair_layout: bool, + destination_catalog_nonoverlapping: bool, + destination_indices_unique: bool, +) -> BackendIndexSelection: + return BackendIndexSelection( + op, + plan_handle, + local_indices, + remote_indices, + selected_nbytes=selected_nbytes, + uniform_pair_layout=uniform_pair_layout, + destination_catalog_nonoverlapping=destination_catalog_nonoverlapping, + destination_indices_unique=destination_indices_unique, + _token=_INDEX_SELECTION_TOKEN, + ) + + +@dataclass(frozen=True, slots=True) +class RegisteredMemory: + """A local allocation whose lifetime is retained by PyTorch Core.""" + + registration_id: str + name: str + owner: object = field(repr=False) + address: int + nbytes: int + device: str + memory_type: str + + +@dataclass(frozen=True, slots=True) +class ImportedRegistration: + """A named remote allocation imported from endpoint metadata.""" + + peer_handle: object = field(repr=False) + registration_id: str + name: str + nbytes: int + device: str + memory_type: str + descriptor: bytes = field(repr=False) + + +@dataclass(frozen=True, slots=True) +class ImportedRawMemory: + """An expert-only remote span interpreted and validated by a backend.""" + + peer_handle: object = field(repr=False) + address: int + nbytes: int + device: str + memory_type: str + + +class BackendWork(ABC): + """Provider request used by the stable :class:`Work` wrapper. + + ``close`` must be idempotent. Core retains every submission argument until + the public Work is closed. Core also retains any CUDA readiness token even + though it resolves readiness before calling the provider. State may remain + PENDING, advance to RUNNING, or advance directly to a terminal state; + RUNNING may only remain RUNNING or advance to a terminal state. + """ + + # An empty base layout lets latency-sensitive provider subclasses opt into + # real slots. A subclass that omits ``__slots__`` still receives its normal + # instance dictionary, preserving the source-compatible extension model. + __slots__ = () + + @property + @abstractmethod + def state(self) -> WorkState | None: ... + + @property + @abstractmethod + def error(self) -> BaseException | None: ... + + @property + def completion_event(self) -> object | None: + return None + + def poll_bounded( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState: + """Perform at most ``max_polls`` state observations. + + Providers advertising ``capabilities.fused_work_poll`` should override + this method so the bounded loop remains inside one provider boundary. + The default preserves source compatibility with existing providers. + A zero timeout still permits the first observation; the timeout is + checked only between observations and cannot preempt a blocking poll. + """ + + if isinstance(max_polls, bool) or not isinstance(max_polls, int): + raise TypeError("max_polls must be a positive integer") + if max_polls <= 0 or max_polls > _SIGNED_INT64_MAX: + raise ValueError("max_polls must be a positive signed 64-bit integer") + if timeout_ns is not None: + if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): + raise TypeError("timeout_ns must be a non-negative integer or None") + if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + if isinstance(timeout_check_interval, bool) or not isinstance( + timeout_check_interval, int + ): + raise TypeError("timeout_check_interval must be a positive integer") + if timeout_check_interval <= 0 or timeout_check_interval > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_check_interval must be a positive signed 64-bit integer" + ) + deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns + state = WorkState.PENDING + for poll_index in range(max_polls): + state = self.state + if ( + not isinstance(state, WorkState) + or state in _REQUEST_SLOT_TERMINAL_STATES + ): + break + if ( + deadline_ns is not None + and (poll_index == 0 or (poll_index + 1) % timeout_check_interval == 0) + and time.monotonic_ns() >= deadline_ns + ): + break + return state + + @abstractmethod + def cancel(self) -> bool: + """Request best-effort cancellation. + + Returns ``True`` only when the request is known to be cancelled. + """ + + @abstractmethod + def close(self) -> None: + """Release the provider request after it reaches a terminal state.""" + + +class BackendRequestSlot(ABC): + """Provider-owned persistent request behind a public ``TransferSlot``. + + Core queries ``state``, ``error``, and ``completion_event`` only while one + unrecycled generation is active. ``start`` rearms the immutable transfer + described at construction. It must return ``None`` and, once native work + may have launched, must convert every launch failure (including a caught + ``BaseException``) into observable terminal slot state rather than raise. + + ``recycle`` is called only after terminal state and returns the request to + provider-idle state without changing its immutable transfer description. + It must return ``None`` and be idempotent after successful recycle so Core + can retry if an asynchronous interruption lands after the provider state + transition. Ordinarily, ``close`` is called only while provider-idle and + must likewise be idempotent after successful release. + + A provider advertising + ``capabilities.deferred_successful_request_slot_recycle`` additionally + retains an observable, clean ``COMPLETED`` receipt when Core omits the + successful recycle in ``TransferSlot.execute``. Its next ``start``, + ``start_and_poll``, ``start_with_notification``, + ``start_and_poll_with_notification``, or ``close`` must consume that + receipt inside the same provider boundary. Explicit ``recycle`` remains + supported and idempotent. + Failure, cancellation, ambiguous launch, and nonterminal state must never + be consumed implicitly, and ``poll_state`` must continue to report + ``COMPLETED`` until the receipt is consumed. + + Core serializes calls to one slot under ``SERIALIZED`` and ``MULTIPLE``; + ``SINGLE`` supplies the same guarantee through its one-owner-thread + contract. Providers therefore need synchronization only for state shared + between different slots or other backend resources. Independent slots may + make progress concurrently under ``MULTIPLE``. Asynchronous interruption + is recoverable at every provider boundary, but request-slot methods are not + recursively reentrant from a Python signal handler while the interrupted + method is still unwinding. + """ + + # See BackendWork: provider slots can be dictionary-free without requiring + # existing unslotted third-party implementations to change. + __slots__ = () + + @property + @abstractmethod + def state(self) -> WorkState | None: ... + + def poll_state(self) -> WorkState | None: + """Return current state through a cacheable method dispatch. + + Providers with a latency-sensitive request path should override this + method directly. The default preserves source compatibility with + providers that implement only the ``state`` property. ``None`` means + the provider is still idle because an interrupted ``start`` never + entered or acknowledged a generation. Once a generation is + acknowledged, the provider must report a ``WorkState`` until recycle. + """ + + return self.state + + @property + @abstractmethod + def error(self) -> BaseException | None: ... + + @property + def completion_event(self) -> object | None: + return None + + @abstractmethod + def start(self) -> None: + """Start one generation of the immutable request.""" + + def start_with_notification(self, notification: bytes | None) -> None: + """Start with an explicit request-local notification override. + + ``None`` explicitly clears the notification for this generation; + non-``None`` payloads are non-empty immutable bytes already validated + by Core. Before returning, the provider must synchronously copy or + retain the payload for as long as native work can access it; Core need + not retain the argument after this boundary. Providers advertise this + optional operation through + ``request_slot_notification_overrides``. The default is deliberately + unsupported so an older provider can never silently retain a stale + notification from a prior repost. + """ + + del notification + raise UnsupportedError( + "backend request slot does not support notification overrides" + ) + + def start_and_poll( + self, + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + """Start one generation and perform a bounded initial state poll. + + This source-compatible default composes :meth:`start` and + :meth:`poll_state`. Providers advertising + ``capabilities.fused_request_slot_start_poll`` should override it so + native post and status polling remain inside one provider boundary. + + At most ``max_polls`` state observations may be made. Zero is a + post-only budget: fused providers should return the state implied by + the native post result without issuing a status query. The default + source-compatible composition returns the cached ``state`` after + ``start``; only an overriding fused provider can guarantee that reading + it performs no native probe. ``timeout_ns`` bounds elapsed time from + entry, checked between observations; it cannot preempt a blocking + provider operation. With a positive budget, one observation is made + after start even for a zero timeout. A successful implementation must + return a :class:`WorkState`. ``None`` retains its idle/unacknowledged + meaning so Core can recover an interrupted start, but is not a valid + successful acknowledgement of a generation. + """ + + if isinstance(max_polls, bool) or not isinstance(max_polls, int): + raise TypeError("max_polls must be a non-negative integer") + if max_polls < 0 or max_polls > _SIGNED_INT64_MAX: + raise ValueError("max_polls must be a non-negative signed 64-bit integer") + if timeout_ns is not None: + if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): + raise TypeError("timeout_ns must be a non-negative integer or None") + if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + start_result = self.start() + if start_result is not None: + raise TypeError("backend request-slot start must return None") + if max_polls == 0: + return self.state + deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns + state: WorkState | None = None + for _ in range(max_polls): + state = self.poll_state() + if ( + state is None + or not isinstance(state, WorkState) + or state in _REQUEST_SLOT_TERMINAL_STATES + or (deadline_ns is not None and time.monotonic_ns() >= deadline_ns) + ): + break + return state + + def start_and_poll_with_notification( + self, + notification: bytes | None, + *, + max_polls: int, + timeout_ns: int | None, + ) -> WorkState | None: + """Start with a notification override and perform bounded polling. + + This source-compatible default composes + :meth:`start_with_notification` and :meth:`poll_state`. Providers that + also advertise ``fused_request_slot_start_poll`` should override it so + notification selection, native post, and the initial status probes + remain in one provider boundary. + + The same copy-or-retain requirement as :meth:`start_with_notification` + applies before this method returns. A zero poll budget has the same + post-only semantics and default limitation as :meth:`start_and_poll`. + """ + + if isinstance(max_polls, bool) or not isinstance(max_polls, int): + raise TypeError("max_polls must be a non-negative integer") + if max_polls < 0 or max_polls > _SIGNED_INT64_MAX: + raise ValueError("max_polls must be a non-negative signed 64-bit integer") + if timeout_ns is not None: + if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): + raise TypeError("timeout_ns must be a non-negative integer or None") + if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + result = self.start_with_notification(notification) + if result is not None: + raise TypeError( + "backend request-slot start_with_notification must return None" + ) + if max_polls == 0: + return self.state + deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns + state: WorkState | None = None + for _ in range(max_polls): + state = self.poll_state() + if ( + state is None + or not isinstance(state, WorkState) + or state in _REQUEST_SLOT_TERMINAL_STATES + or (deadline_ns is not None and time.monotonic_ns() >= deadline_ns) + ): + break + return state + + def poll_bounded( + self, + *, + max_polls: int, + timeout_ns: int | None, + timeout_check_interval: int, + ) -> WorkState | None: + """Perform bounded state observations for an active generation. + + Providers advertising ``capabilities.fused_request_slot_poll`` should + override this source-compatible default so multiple native status + observations remain inside one provider boundary. The capability does + not require start/poll fusion. + + At most ``max_polls`` observations are made. ``timeout_ns`` measures a + soft monotonic elapsed-time bound from entry; it is checked after the + first observation and then every ``timeout_check_interval`` + observations. It cannot preempt a blocking provider call, and sampled + checking may overshoot by up to one sampling interval of provider + calls. A zero timeout still makes exactly one observation. ``None`` + disables clock checks. A returned ``None`` retains its provider-idle + meaning and is invalid after Core has observed an acknowledged active + generation. + """ + + if isinstance(max_polls, bool) or not isinstance(max_polls, int): + raise TypeError("max_polls must be a positive integer") + if max_polls <= 0 or max_polls > _SIGNED_INT64_MAX: + raise ValueError("max_polls must be a positive signed 64-bit integer") + if timeout_ns is not None: + if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): + raise TypeError("timeout_ns must be a non-negative integer or None") + if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_ns must be a non-negative signed 64-bit integer or None" + ) + if isinstance(timeout_check_interval, bool) or not isinstance( + timeout_check_interval, int + ): + raise TypeError("timeout_check_interval must be a positive integer") + if timeout_check_interval <= 0 or timeout_check_interval > _SIGNED_INT64_MAX: + raise ValueError( + "timeout_check_interval must be a positive signed 64-bit integer" + ) + deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns + state: WorkState | None = None + for poll in range(max_polls): + state = self.poll_state() + if ( + state is None + or not isinstance(state, WorkState) + or state in _REQUEST_SLOT_TERMINAL_STATES + ): + break + observations = poll + 1 + if ( + deadline_ns is not None + and (observations == 1 or observations % timeout_check_interval == 0) + and time.monotonic_ns() >= deadline_ns + ): + break + return state + + @abstractmethod + def cancel(self) -> bool: + """Request best-effort cancellation of the active generation.""" + + @abstractmethod + def recycle(self) -> None: + """Return a terminal generation to retry-safe provider-idle state.""" + + @abstractmethod + def close(self) -> None: + """Release this provider request while it is idle.""" + + +class Backend(ABC): + """Experimental provider SPI for one-sided endpoint transfer. + + Background progress is provider-owned. Release methods must be idempotent + after a successful release and must make partial cleanup retry-safe. + Handle-producing methods construct a provider transaction record, pass it + to ``adopt_handle`` before their first native mutation, and return ``None``. + Core can therefore roll back an adopted partial mutation or retain it for a + public close retry. Before adoption, a provider must fully roll back any + internal state before raising. If native cleanup fails, it may additionally + retain the orphan internally, reject conflicting operations, and retry + cleanup on later calls and close. + """ + + @property + @abstractmethod + def capabilities(self) -> Capabilities: ... + + @property + def supported_thread_modes(self) -> frozenset[ThreadMode]: + return frozenset({ThreadMode.SERIALIZED}) + + @abstractmethod + def endpoint_payload(self, registration_handles: Sequence[object]) -> bytes: + """Export connection metadata limited to the selected registrations. + + A provider must not silently grant access to omitted registrations. It + may reject selective export if its transport cannot enforce it. + """ + + ... + + @abstractmethod + def register( + self, + memory: RegisteredMemory, + *, + adopt_handle: Callable[[object], None], + ) -> None: + """Adopt a retry-safe provider record before native registration.""" + + @abstractmethod + def deregister(self, registration_handle: object) -> None: ... + + @abstractmethod + def describe_registration(self, registration_handle: object) -> bytes: ... + + @abstractmethod + def import_peer( + self, + metadata: EndpointMetadata, + *, + adopt_handle: Callable[[object], None], + ) -> None: + """Adopt a retry-safe provider record before native peer mutation.""" + + @abstractmethod + def release_peer(self, peer_handle: object) -> None: ... + + @abstractmethod + def prepare( + self, + local: Sequence[BackendRegion], + remote: Sequence[BackendRegion], + *, + indexed: bool, + adopt_handle: Callable[[object], None], + ) -> None: + """Adopt a retry-safe provider record before native preparation.""" + + @abstractmethod + def release_plan(self, plan_handle: object) -> None: ... + + @abstractmethod + def submit( + self, + op: TransferOp, + plan_handle: object, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | None, + reuse_token: object | None, + adopt_work: Callable[[BackendWork], None], + ) -> None: + """Submit work after Core has established local CUDA readiness. + + Core resolves the selected CUDA ordering mode before entering the + serialized provider call, so this method must not add an implicit device + synchronization. Terminal completion means local memory may be safely + reused. ``completion_event``, when present, must express that completion + to the device. An attached notification may become visible only after the + remote writes are visible. + + Providers must also reject overlapping physical destination spans, + including aliases advertised as distinct remote registrations. + + For an indexed plan, each prepared ``BackendRegion`` contributes + ``count`` logical blocks, or one block when ``count`` is absent. Indices + address the flattened region-major, block-minor catalogs. A logical + block has ``nbytes`` bytes at ``offset + within_region_index * stride`` + (with an effective stride of ``nbytes`` for an unstrided region); + providers may keep this representation compact rather than expanding it + at prepare. ``None`` indices select every logical block in that catalog. + + ``reuse_token`` is an opaque identity hint. Core passes the same + non-``None`` object for every submission from one ``BoundTransfer`` and + for repeated full-selection, notification-free, caller-ready convenience + submissions of one operation on one plan. A provider may key immutable + cached lowering state by object identity, but must not inspect the token, + derive transfer semantics from it, or require it to be non-``None``. + Core tokens support weak references; caches should be weak-keyed or + scoped to and cleared with the provider plan. + + ``submit_prevalidated`` has a narrower token contract: its token + identifies the exact plan/operation/index selection, while notification + is request-local and may change. Only an override that can update the + native notification at post time may reuse that token. The default + therefore clears it when forwarding to this legacy method. + + Submission is transactional with respect to ownership. The provider + must construct its BackendWork and synchronously call ``adopt_work`` + before launching transport work. Before adoption it may raise only if + no transport access can remain in flight. After successful adoption it + must not raise: launch failures, including asynchronous interruptions, + must instead become observable through the adopted BackendWork. This + handshake ensures Core owns all memory and plan lifetimes before native + code can access them. If the adoption callback itself is interrupted, + the provider must not launch; it should make its Work terminally failed + and return normally, allowing Core to determine whether adoption took + effect before the interruption. + """ + + ... + + def submit_prevalidated( + self, + op: TransferOp, + plan_handle: object, + *, + selection: BackendIndexSelection, + notification: bytes | None, + reuse_token: object | None, + adopt_work: Callable[[BackendWork], None], + ) -> None: + """Submit a Core-authenticated immutable indexed selection. + + The default preserves provider source compatibility by forwarding the + two read-only native-int32 views to :meth:`submit`. A provider may + override this method to consume the packed views without rebuilding + them and to combine Core's proof flags with its own plan-resolution + proofs. It must not trust a duck-typed object. + """ + + if type(selection) is not BackendIndexSelection: + raise TypeError("selection must be a Core BackendIndexSelection") + selection._validate_binding(op, plan_handle) + del reuse_token + return self.submit( + op, + plan_handle, + local_indices=selection.local_indices, + remote_indices=selection.remote_indices, + notification=notification, + reuse_token=None, + adopt_work=adopt_work, + ) + + def submit_prevalidated_and_poll( + self, + op: TransferOp, + plan_handle: object, + *, + selection: BackendIndexSelection, + notification: bytes | None, + max_polls: int, + timeout_ns: int | None, + adopt_work: Callable[[BackendWork], None], + ) -> WorkState: + """Submit one authenticated selection and perform a bounded probe. + + This is capability-gated rather than composed through a callback + wrapper: only a provider that can keep preparation, post, and polling + in its intended native boundary should advertise + ``fused_prevalidated_submit_poll``. Existing providers remain source + compatible and use :meth:`submit_prevalidated` plus public Work polls. + """ + + del ( + op, + plan_handle, + selection, + notification, + max_polls, + timeout_ns, + adopt_work, + ) + raise UnsupportedError( + "backend does not support fused prevalidated submit and poll" + ) + + def create_request_slot( + self, + op: TransferOp, + plan_handle: object, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | None, + adopt_slot: Callable[..., None], + ) -> None: + """Create an opt-in persistent request for one immutable transfer. + + The provider must construct a retry-safe :class:`BackendRequestSlot` + record and synchronously call ``adopt_slot`` before its first native + mutation. A provider may pass an optional keyword-only + ``execution_dispatch`` object to that callback. Core then validates and + caches the dispatch object's execution methods while retaining the + :class:`BackendRequestSlot` wrapper for diagnostics, completion events, + and close. Omitting the keyword keeps the wrapper as the dispatch + object and is source compatible with existing providers. + + Before adoption the provider may raise only after fully rolling back + provider state. If the adoption callback is interrupted, it must not + create or launch native work and should return normally when Core may + already have retained the slot. + + A provider advertises this optional SPI with + ``capabilities.persistent_request_slots``. The default rejects it so + existing providers remain source compatible. + """ + + del op, plan_handle, local_indices, remote_indices, notification, adopt_slot + raise UnsupportedError("backend does not support persistent request slots") + + @abstractmethod + def send_notification( + self, + peer_handle: object, + payload: bytes, + *, + adopt_work: Callable[[BackendWork], None], + ) -> None: + """Send a standalone notification. + + Terminal Work completion has the meaning declared by + ``capabilities.standalone_notification_completion``. + The provider must use the same adopt-before-launch transaction described + by :meth:`submit`. + """ + + ... + + def send_notification_sync(self, peer_handle: object, payload: bytes) -> None: + """Send without creating a provider or Core Work. + + Return means the completion point declared by + ``capabilities.standalone_notification_completion``; it does not imply + remote delivery unless that declaration is ``REMOTE_DELIVERED``. + Providers opt in with ``capabilities.synchronous_notification_send``. + Under ``MULTIPLE``, Core pins the peer lifetime but does not serialize + independent native calls; providers must preserve their advertised + concurrent-entry contract. + """ + + del peer_handle, payload + raise UnsupportedError( + "backend does not support synchronous standalone notification send" + ) + + @abstractmethod + def poll_notifications( + self, max_items: int | None = None + ) -> list[BackendNotification]: + """Destructively receive up to ``max_items`` notifications. + + This v0 operation is at-most-once: the SPI has no batch token or + acknowledge/rollback phase. Providers must validate native ingress + before exposing it, while applications that require reliable delivery + must sequence, acknowledge, deduplicate, and retransmit messages. + """ + + ... + + def poll_notification_batches( + self, max_items: int | None = None + ) -> list[BackendNotificationBatch]: + """Destructively receive notifications while preserving source groups. + + ``max_items`` counts payloads, not batches. Providers with a native + grouped receive should override this method and advertise + ``capabilities.grouped_notification_receive``. This compatibility + default groups only adjacent notifications with a complete matching + identity. Anonymous notifications remain singleton batches because + scalar input cannot prove that they share a native source. + """ + + notifications = self.poll_notifications(max_items) + batches: list[BackendNotificationBatch] = [] + source: tuple[str, str] | None = None + payloads: list[bytes] = [] + + def flush() -> None: + nonlocal source, payloads + if not payloads or source is None: + return + batches.append( + BackendNotificationBatch( + payloads=tuple(payloads), + source_endpoint_id=source[0], + source_incarnation=source[1], + ) + ) + source = None + payloads = [] + + for notification in notifications: + if not isinstance(notification, BackendNotification): + raise TypeError( + "backend notifications must contain BackendNotification values" + ) + if notification.source_endpoint_id is None: + flush() + batches.append(BackendNotificationBatch((notification.payload,))) + continue + assert notification.source_incarnation is not None + next_source = ( + notification.source_endpoint_id, + notification.source_incarnation, + ) + if source != next_source: + flush() + source = next_source + payloads.append(notification.payload) + flush() + return batches + + @abstractmethod + def progress(self) -> None: ... + + @abstractmethod + def close(self) -> None: + """Close the provider; this must be idempotent after native success.""" + + ... + + +BackendFactory = Callable[ + [ + str, + str, + str, + ProgressMode, + ThreadMode, + Mapping[str, Any], + ], + Backend, +] + + +class BackendFactoryV2(Protocol): + """Factory that transfers backend ownership before native activation. + + Version-2 providers must construct a pure, closeable Python backend record, + pass it exactly once to ``adopt_backend``, and only then activate or publish + native state. A successful call returns ``None``. Core deliberately does + not inspect callable signatures or retry a factory under another contract. + """ + + def __call__( + self, + endpoint_id: str, + name: str, + incarnation: str, + progress_mode: ProgressMode, + thread_mode: ThreadMode, + options: Mapping[str, Any], + *, + adopt_backend: Callable[[Backend], None], + ) -> None: ... + + +# This is the maximum backward-compatible factory contract understood by this +# Core. A provider selecting version N may run when this value is >= N. +BACKEND_FACTORY_API_VERSION = 2 + + +@dataclass(frozen=True, slots=True) +class _BackendRegistration: + factory: BackendFactory | BackendFactoryV2 + factory_api_version: int + + +_BACKENDS: dict[str, _BackendRegistration] = {} +_BACKENDS_LOCK = threading.Lock() + + +def register_backend( + name: str, + factory: BackendFactory | BackendFactoryV2, + *, + replace: bool = False, + factory_api_version: int = 1, +) -> None: + """Register an endpoint-transfer backend factory. + + Providers should call this explicitly from an integration module. Importing + :mod:`torch` never imports an external provider. Version 1 is retained for + source compatibility and returns a :class:`Backend`; its return boundary + cannot close the final CALL-to-adoption interruption window. Version 2 + receives an ``adopt_backend`` keyword callback, must call it exactly once + before native activation, and returns ``None``. Repeating a registration + with the same factory object and exact API version is idempotent. This + closes the caller's post-commit/lost-return window without weakening name + collision protection: a different factory object or version remains busy. + """ + + if not isinstance(name, str) or not name: + raise ValueError("backend name must be a non-empty string") + if not callable(factory): + raise TypeError("backend factory must be callable") + if type(factory_api_version) is not int: + raise TypeError("factory_api_version must be an integer") + if factory_api_version not in (1, 2): + raise ValueError("factory_api_version must be 1 or 2") + with _BACKENDS_LOCK: + existing = _BACKENDS.get(name) + if existing is not None and not replace: + if ( + existing.factory is factory + and existing.factory_api_version == factory_api_version + ): + return + raise BusyError(f"transfer backend {name!r} is already registered") + _BACKENDS[name] = _BackendRegistration(factory, factory_api_version) + + +def unregister_backend(name: str) -> None: + """Remove a provider factory. + + This only affects future endpoints and is primarily intended for tests. + Existing endpoints continue to own their backend instances. + """ + + with _BACKENDS_LOCK: + if name not in _BACKENDS: + raise BackendNotFoundError(f"transfer backend {name!r} is not registered") + del _BACKENDS[name] + + +def _get_backend_registration(name: str) -> _BackendRegistration: + with _BACKENDS_LOCK: + try: + return _BACKENDS[name] + except KeyError as error: + raise BackendNotFoundError( + f"transfer backend {name!r} is not registered" + ) from error + + +def get_backend_factory(name: str) -> BackendFactory | BackendFactoryV2: + """Return the registered callable without changing the legacy query API.""" + + return _get_backend_registration(name).factory diff --git a/torch/distributed/_transfer/_errors.py b/torch/distributed/_transfer/_errors.py new file mode 100644 index 0000000000000..0dd44bea6d4d6 --- /dev/null +++ b/torch/distributed/_transfer/_errors.py @@ -0,0 +1,62 @@ +__all__ = [ + "BackendError", + "BackendFailureError", + "BackendMismatchError", + "BackendNotFoundError", + "BusyError", + "ClosedError", + "InvalidMetadataError", + "InvalidRegionError", + "TransferCancelledError", + "TransferError", + "TransferTimeoutError", + "UnsupportedError", +] + + +class TransferError(RuntimeError): + """Base exception for endpoint transfer failures.""" + + +class BackendError(TransferError): + """Base exception for transfer backend failures.""" + + +class BackendNotFoundError(BackendError): + """Raised when a requested backend is not registered.""" + + +class BackendMismatchError(BackendError): + """Raised when metadata belongs to a different backend.""" + + +class BackendFailureError(BackendError): + """Raised when a backend operation fails.""" + + +class InvalidMetadataError(TransferError): + """Raised when endpoint metadata is malformed or unsupported.""" + + +class InvalidRegionError(TransferError): + """Raised when a memory region is invalid.""" + + +class UnsupportedError(TransferError): + """Raised when a requested operation is unsupported.""" + + +class ClosedError(TransferError): + """Raised when an operation targets a closed resource.""" + + +class BusyError(TransferError): + """Raised when a resource cannot be changed while work is in flight.""" + + +class TransferTimeoutError(TransferError): + """Raised when transfer work exceeds its deadline.""" + + +class TransferCancelledError(TransferError): + """Raised when transfer work is cancelled.""" diff --git a/torch/distributed/_transfer/_reference.py b/torch/distributed/_transfer/_reference.py new file mode 100644 index 0000000000000..4f898c08ff38f --- /dev/null +++ b/torch/distributed/_transfer/_reference.py @@ -0,0 +1,724 @@ +from __future__ import annotations + +import ctypes +import itertools +import json +import threading +from collections import deque +from dataclasses import dataclass +from typing import TYPE_CHECKING +from weakref import WeakValueDictionary + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping, Sequence + from typing import Any + +from ._backend import ( + Backend, + BackendWork, + ImportedRawMemory, + ImportedRegistration, + RegisteredMemory, + register_backend, +) +from ._errors import ( + BusyError, + ClosedError, + InvalidMetadataError, + InvalidRegionError, + UnsupportedError, +) +from ._types import ( + BackendNotification, + BackendRegion, + Capabilities, + EndpointMetadata, + NotificationCompletion, + ProgressMode, + ThreadMode, + TransferOp, + WorkState, +) + + +@dataclass(eq=False, slots=True) +class _ReferencePeer: + backend: _ReferenceBackend + endpoint_id: str + incarnation: str + allowed_registration_ids: frozenset[str] + released: bool = False + + +@dataclass(eq=False, slots=True) +class _ReferencePlan: + local: tuple[BackendRegion, ...] + remote: tuple[BackendRegion, ...] + indexed: bool + released: bool = False + + +class _ReferenceWork(BackendWork): + def __init__( + self, + state: WorkState, + error: BaseException | None = None, + ) -> None: + self._state = state + self._error = error + self._closed = False + + @property + def state(self) -> WorkState: + if self._closed: + raise ClosedError("reference work is closed") + return self._state + + @property + def error(self) -> BaseException | None: + if self._closed: + raise ClosedError("reference work is closed") + return self._error + + def cancel(self) -> bool: + if self._closed: + raise ClosedError("reference work is closed") + if self._state in { + WorkState.COMPLETED, + WorkState.FAILED, + WorkState.CANCELLED, + }: + return self._state is WorkState.CANCELLED + self._state = WorkState.CANCELLED + return True + + def close(self) -> None: + self._closed = True + + +class _ReferenceBackend(Backend): + """Synchronous, in-process CPU backend used for contract testing.""" + + _CAPABILITIES = Capabilities( + read=True, + write=True, + attached_notifications=True, + standalone_notifications=True, + raw_spans=False, + strided_regions=True, + manual_progress=True, + background_progress=True, + synchronous_notification_send=True, + standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, + ) + + def __init__( + self, + endpoint_id: str, + name: str, + incarnation: str, + progress_mode: ProgressMode, + thread_mode: ThreadMode, + options: Mapping[str, Any], + ) -> None: + self._initialize_record( + endpoint_id, + name, + incarnation, + progress_mode, + thread_mode, + options, + ) + try: + self._activate() + except BaseException as error: + try: + self.close() + except BaseException as cleanup_error: + add_note = getattr(error, "add_note", None) + if callable(add_note): + add_note( + "reference backend activation cleanup also raised: " + f"{cleanup_error}" + ) + raise + + @classmethod + def _new_unactivated( + cls, + endpoint_id: str, + name: str, + incarnation: str, + progress_mode: ProgressMode, + thread_mode: ThreadMode, + options: Mapping[str, Any], + ) -> _ReferenceBackend: + backend = cls.__new__(cls) + backend._initialize_record( + endpoint_id, + name, + incarnation, + progress_mode, + thread_mode, + options, + ) + return backend + + def _initialize_record( + self, + endpoint_id: str, + name: str, + incarnation: str, + progress_mode: ProgressMode, + thread_mode: ThreadMode, + options: Mapping[str, Any], + ) -> None: + del progress_mode, thread_mode + if options: + unknown = ", ".join(sorted(options)) + raise ValueError(f"reference backend has no options: {unknown}") + self._endpoint_id = endpoint_id + self._name = name + self._incarnation = incarnation + self._registrations: dict[str, RegisteredMemory] = {} + self._plans: set[_ReferencePlan] = set() + self._peers: set[_ReferencePeer] = set() + self._notifications: deque[BackendNotification] = deque() + self._lock = threading.RLock() + self._activation_started = False + self._activated = False + self._closed = False + + def _activate(self) -> None: + if self._closed: + raise ClosedError("cannot activate a closed reference backend") + if self._activated: + return + if self._activation_started: + raise BusyError("reference backend activation is already in progress") + self._activation_started = True + with _REFERENCE_ENDPOINTS_LOCK: + if self._endpoint_id in _REFERENCE_ENDPOINTS: + raise BusyError(f"duplicate reference endpoint {self._endpoint_id!r}") + _REFERENCE_ENDPOINTS[self._endpoint_id] = self + self._activated = True + + @property + def capabilities(self) -> Capabilities: + return self._CAPABILITIES + + @property + def supported_thread_modes(self) -> frozenset[ThreadMode]: + return frozenset( + { + ThreadMode.SINGLE, + ThreadMode.CALLER_SERIALIZED, + ThreadMode.SERIALIZED, + ThreadMode.MULTIPLE, + } + ) + + def endpoint_payload(self, registration_handles: Sequence[object]) -> bytes: + self._check_open() + registration_ids = [ + _require_registration_handle(handle) for handle in registration_handles + ] + with self._lock: + if any( + registration_id not in self._registrations + for registration_id in registration_ids + ): + raise ClosedError("reference registration is not active") + return json.dumps( + { + "endpoint_id": self._endpoint_id, + "incarnation": self._incarnation, + "registrations": sorted(registration_ids), + }, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + def register( + self, + memory: RegisteredMemory, + *, + adopt_handle: Callable[[object], None], + ) -> None: + self._check_open() + if not memory.device.startswith("cpu"): + raise UnsupportedError("reference backend supports CPU allocations only") + with self._lock: + if memory.registration_id in self._registrations: + raise BusyError("duplicate reference registration") + adopt_handle(memory.registration_id) + self._registrations[memory.registration_id] = memory + + def deregister(self, registration_handle: object) -> None: + self._check_open() + registration_id = _require_registration_handle(registration_handle) + with self._lock: + self._registrations.pop(registration_id, None) + + def describe_registration(self, registration_handle: object) -> bytes: + self._check_open() + registration_id = _require_registration_handle(registration_handle) + with self._lock: + if registration_id not in self._registrations: + raise ClosedError("reference registration is not active") + return registration_id.encode("utf-8") + + def import_peer( + self, + metadata: EndpointMetadata, + *, + adopt_handle: Callable[[object], None], + ) -> None: + self._check_open() + with _REFERENCE_ENDPOINTS_LOCK: + try: + remote = _REFERENCE_ENDPOINTS[metadata.endpoint_id] + except KeyError as error: + raise InvalidMetadataError( + "reference peer is not present in this process" + ) from error + if remote._incarnation != metadata.incarnation: + raise InvalidMetadataError("reference peer incarnation is stale") + try: + payload = json.loads(metadata.backend_payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise InvalidMetadataError("reference peer payload is malformed") from error + if not isinstance(payload, dict) or set(payload) != { + "endpoint_id", + "incarnation", + "registrations", + }: + raise InvalidMetadataError("reference peer payload has invalid fields") + registration_ids = payload["registrations"] + if ( + payload["endpoint_id"] != metadata.endpoint_id + or payload["incarnation"] != metadata.incarnation + or not isinstance(registration_ids, list) + or not all(isinstance(value, str) for value in registration_ids) + or len(registration_ids) != len(set(registration_ids)) + ): + raise InvalidMetadataError("reference peer payload does not match") + advertised_ids = {item.registration_id for item in metadata.registrations} + if set(registration_ids) != advertised_ids: + raise InvalidMetadataError( + "reference payload registration grant does not match metadata" + ) + peer = _ReferencePeer( + backend=remote, + endpoint_id=metadata.endpoint_id, + incarnation=metadata.incarnation, + allowed_registration_ids=frozenset(registration_ids), + ) + adopt_handle(peer) + self._peers.add(peer) + + def release_peer(self, peer_handle: object) -> None: + self._check_open() + peer = _require_peer(peer_handle) + if peer.released: + return + peer.released = True + self._peers.discard(peer) + + def prepare( + self, + local: Sequence[BackendRegion], + remote: Sequence[BackendRegion], + *, + indexed: bool, + adopt_handle: Callable[[object], None], + ) -> None: + self._check_open() + for region in local: + if region.remote: + raise InvalidRegionError("local reference region is remote") + self._resolve_local(region) + for region in remote: + if not region.remote: + raise InvalidRegionError("remote reference region is local") + self._resolve_remote(region) + plan = _ReferencePlan(tuple(local), tuple(remote), indexed) + adopt_handle(plan) + self._plans.add(plan) + + def release_plan(self, plan_handle: object) -> None: + self._check_open() + plan = _require_plan(plan_handle) + if plan.released: + return + plan.released = True + self._plans.discard(plan) + + def submit( + self, + op: TransferOp, + plan_handle: object, + *, + local_indices: Sequence[int] | None, + remote_indices: Sequence[int] | None, + notification: bytes | None, + reuse_token: object | None, + adopt_work: Callable[[BackendWork], None], + ) -> None: + self._check_open() + plan = _require_plan(plan_handle) + if plan.released: + raise ClosedError("reference plan is released") + work = _ReferenceWork(WorkState.PENDING) + try: + adopt_work(work) + except BaseException as error: + # A callback interruption can occur immediately before or after + # Core stores the reference. Never launch in either case; if Core + # did adopt it, expose a terminal failure instead of re-raising. + work._error = error + work._state = WorkState.FAILED + return + try: + local = _select(plan.local, local_indices, indexed=plan.indexed) + remote = _select(plan.remote, remote_indices, indexed=plan.indexed) + local_spans = [ + span for region in local for span in self._resolve_local(region) + ] + remote_memory = [self._resolve_remote(region) for region in remote] + remote_spans = [span for spans, _ in remote_memory for span in spans] + if op is TransferOp.WRITE: + source, destination = local_spans, remote_spans + elif op is TransferOp.READ: + source, destination = remote_spans, local_spans + else: + raise UnsupportedError(f"unsupported transfer operation {op!r}") + _copy_spans( + source, + destination, + anchors=tuple(memory for _, memory in remote_memory), + ) + if notification is not None: + peer = self._plan_peer(plan) + peer.backend._enqueue_notification( + BackendNotification( + payload=notification, + source_endpoint_id=self._endpoint_id, + source_incarnation=self._incarnation, + ) + ) + work._state = WorkState.COMPLETED + except BaseException as error: + work._error = error + work._state = WorkState.FAILED + + def send_notification( + self, + peer_handle: object, + payload: bytes, + *, + adopt_work: Callable[[BackendWork], None], + ) -> None: + self._check_open() + work = _ReferenceWork(WorkState.PENDING) + try: + adopt_work(work) + except BaseException as error: + work._error = error + work._state = WorkState.FAILED + return + try: + peer = _require_peer(peer_handle) + if peer.released: + raise ClosedError("reference peer is released") + peer.backend._enqueue_notification( + BackendNotification( + payload=payload, + source_endpoint_id=self._endpoint_id, + source_incarnation=self._incarnation, + ) + ) + work._state = WorkState.COMPLETED + except BaseException as error: + work._error = error + work._state = WorkState.FAILED + + def send_notification_sync(self, peer_handle: object, payload: bytes) -> None: + self._check_open() + peer = _require_peer(peer_handle) + if peer.released: + raise ClosedError("reference peer is released") + peer.backend._enqueue_notification( + BackendNotification( + payload=payload, + source_endpoint_id=self._endpoint_id, + source_incarnation=self._incarnation, + ) + ) + + def poll_notifications( + self, max_items: int | None = None + ) -> list[BackendNotification]: + self._check_open() + count = len(self._notifications) if max_items is None else max_items + result: list[BackendNotification] = [] + with self._lock: + for _ in range(min(count, len(self._notifications))): + result.append(self._notifications.popleft()) + return result + + def progress(self) -> None: + self._check_open() + + def close(self) -> None: + with self._lock: + if self._closed: + return + if self._registrations or self._plans or self._peers: + raise BusyError("reference backend still owns child resources") + with _REFERENCE_ENDPOINTS_LOCK: + if _REFERENCE_ENDPOINTS.get(self._endpoint_id) is self: + _REFERENCE_ENDPOINTS.pop(self._endpoint_id, None) + self._notifications.clear() + self._activated = False + self._closed = True + + def _resolve_local(self, region: BackendRegion) -> list[tuple[int, int]]: + registration_id = _require_registration_handle(region.backend_descriptor) + with self._lock: + try: + memory = self._registrations[registration_id] + except KeyError as error: + raise ClosedError( + "local reference registration is not active" + ) from error + return _memory_spans(memory, region) + + def _resolve_remote( + self, region: BackendRegion + ) -> tuple[list[tuple[int, int]], RegisteredMemory]: + descriptor = region.backend_descriptor + if isinstance(descriptor, ImportedRawMemory): + raise UnsupportedError("reference backend rejects remote raw spans") + if not isinstance(descriptor, ImportedRegistration): + raise InvalidRegionError("invalid remote registration descriptor") + peer = _require_peer(descriptor.peer_handle) + if peer.released: + raise ClosedError("reference peer is released") + if descriptor.descriptor != descriptor.registration_id.encode("utf-8"): + raise InvalidMetadataError( + "reference registration descriptor does not match its id" + ) + if descriptor.registration_id not in peer.allowed_registration_ids: + raise InvalidMetadataError( + "remote registration is outside the imported grant" + ) + with peer.backend._lock: + try: + memory = peer.backend._registrations[descriptor.registration_id] + except KeyError as error: + raise ClosedError( + "remote reference registration is not active" + ) from error + if ( + memory.name != descriptor.name + or memory.nbytes != descriptor.nbytes + or memory.device != descriptor.device + or memory.memory_type != descriptor.memory_type + ): + raise InvalidMetadataError("remote registration attributes changed") + return _memory_spans(memory, region), memory + + def _plan_peer(self, plan: _ReferencePlan) -> _ReferencePeer: + descriptor = plan.remote[0].backend_descriptor + if not isinstance(descriptor, ImportedRegistration): + raise InvalidRegionError("reference plan has no named remote peer") + return _require_peer(descriptor.peer_handle) + + def _enqueue_notification(self, notification: BackendNotification) -> None: + with self._lock: + self._check_open() + self._notifications.append(notification) + + def _check_open(self) -> None: + if self._closed or not self._activated: + raise ClosedError(f"reference endpoint {self._name!r} is closed") + + +def _create_reference_backend( + endpoint_id: str, + name: str, + incarnation: str, + progress_mode: ProgressMode, + thread_mode: ThreadMode, + options: Mapping[str, Any], + *, + adopt_backend: Callable[[Backend], None], +) -> None: + backend = _ReferenceBackend._new_unactivated( + endpoint_id, + name, + incarnation, + progress_mode, + thread_mode, + options, + ) + adopt_backend(backend) + backend._activate() + + +def _register_reference_backend() -> None: + register_backend("reference", _create_reference_backend, factory_api_version=2) + + +def _require_registration_handle(value: object) -> str: + if not isinstance(value, str) or not value: + raise InvalidRegionError("invalid reference registration handle") + return value + + +def _require_peer(value: object) -> _ReferencePeer: + if not isinstance(value, _ReferencePeer): + raise InvalidRegionError("invalid reference peer handle") + return value + + +def _require_plan(value: object) -> _ReferencePlan: + if not isinstance(value, _ReferencePlan): + raise InvalidRegionError("invalid reference plan handle") + return value + + +def _select( + regions: tuple[BackendRegion, ...], + indices: Sequence[int] | None, + *, + indexed: bool, +) -> tuple[BackendRegion, ...]: + if indices is None: + return regions + if not indexed: + raise InvalidRegionError("indices require an indexed reference plan") + return tuple(_logical_block(regions, index) for index in indices) + + +def _logical_block( + regions: Sequence[BackendRegion], logical_index: int +) -> BackendRegion: + if isinstance(logical_index, bool) or not isinstance(logical_index, int): + raise InvalidRegionError("reference block indices must be integers") + if logical_index < 0: + raise InvalidRegionError("reference block index must be non-negative") + first_index = 0 + for region in regions: + count = region.count or 1 + if logical_index < first_index + count: + block_index = logical_index - first_index + stride = region.stride or region.nbytes + return BackendRegion( + owner=region.owner, + offset=region.offset + block_index * stride, + nbytes=region.nbytes, + device=region.device, + memory_type=region.memory_type, + backend_descriptor=region.backend_descriptor, + remote=region.remote, + ) + first_index += count + raise InvalidRegionError("reference block index is outside the region catalog") + + +def _memory_spans( + memory: RegisteredMemory, + region: BackendRegion, +) -> list[tuple[int, int]]: + count = region.count or 1 + stride = region.stride or region.nbytes + end = region.offset + (count - 1) * stride + region.nbytes + if end > memory.nbytes: + raise InvalidRegionError("reference region exceeds registered memory") + return [ + (memory.address + region.offset + index * stride, region.nbytes) + for index in range(count) + ] + + +def _copy_spans( + source: list[tuple[int, int]], + destination: list[tuple[int, int]], + *, + anchors: tuple[RegisteredMemory, ...] = (), +) -> None: + # This argument deliberately retains remote allocation owners across every + # raw pointer access, even if the exporting endpoint deregisters meanwhile. + _ = anchors + source_size = sum(nbytes for _, nbytes in source) + destination_size = sum(nbytes for _, nbytes in destination) + if source_size != destination_size: + raise InvalidRegionError( + f"source and destination sizes differ: {source_size} != {destination_size}" + ) + _reject_overlapping_spans(destination) + for source_address, source_nbytes in source: + for destination_address, destination_nbytes in destination: + if _overlaps( + source_address, + source_nbytes, + destination_address, + destination_nbytes, + ): + raise InvalidRegionError( + "reference backend rejects overlapping source and " + "destination regions" + ) + + source_index = destination_index = 0 + source_offset = destination_offset = 0 + while source_index < len(source): + source_address, source_nbytes = source[source_index] + destination_address, destination_nbytes = destination[destination_index] + chunk = min( + source_nbytes - source_offset, + destination_nbytes - destination_offset, + ) + ctypes.memmove( + destination_address + destination_offset, + source_address + source_offset, + chunk, + ) + source_offset += chunk + destination_offset += chunk + if source_offset == source_nbytes: + source_index += 1 + source_offset = 0 + if destination_offset == destination_nbytes: + destination_index += 1 + destination_offset = 0 + + +def _reject_overlapping_spans(spans: list[tuple[int, int]]) -> None: + intervals = sorted((address, address + nbytes) for address, nbytes in spans) + for previous, current in itertools.pairwise(intervals): + if current[0] < previous[1]: + raise InvalidRegionError( + "reference backend rejects overlapping destination regions" + ) + + +def _overlaps( + first_address: int, + first_nbytes: int, + second_address: int, + second_nbytes: int, +) -> bool: + return ( + first_address < second_address + second_nbytes + and second_address < first_address + first_nbytes + ) + + +_REFERENCE_ENDPOINTS: WeakValueDictionary[str, _ReferenceBackend] = ( + WeakValueDictionary() +) +_REFERENCE_ENDPOINTS_LOCK = threading.Lock() diff --git a/torch/distributed/_transfer/_types.py b/torch/distributed/_transfer/_types.py new file mode 100644 index 0000000000000..f696d01f1db67 --- /dev/null +++ b/torch/distributed/_transfer/_types.py @@ -0,0 +1,867 @@ +import base64 +import binascii +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, ClassVar + +from ._errors import InvalidMetadataError, InvalidRegionError + + +__all__ = [ + "BackendNotification", + "BackendNotificationBatch", + "BackendRegion", + "Capabilities", + "CudaOrderingMode", + "EndpointIdentity", + "EndpointInfo", + "EndpointMetadata", + "Notification", + "NotificationBatch", + "NotificationCompletion", + "ProgressMode", + "RegistrationMetadata", + "ThreadMode", + "TransferOp", + "WorkState", +] + + +class WorkState(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class TransferOp(str, Enum): + READ = "read" + WRITE = "write" + + +class ProgressMode(str, Enum): + MANUAL = "manual" + BACKGROUND = "background" + + +class ThreadMode(str, Enum): + """Concurrency contract for one Endpoint resource graph. + + ``SINGLE`` is confined to the creating thread. ``CALLER_SERIALIZED`` may + migrate between threads, but the caller must prevent overlapping or + reentrant operations across the endpoint and every derived resource. + ``SERIALIZED`` asks Core to provide that exclusion. ``MULTIPLE`` permits + concurrent calls and requires a provider that supports them. + """ + + SINGLE = "single" + CALLER_SERIALIZED = "caller_serialized" + SERIALIZED = "serialized" + MULTIPLE = "multiple" + + +class CudaOrderingMode(str, Enum): + """How transport access is ordered after CUDA producer work.""" + + CALLER_READY = "caller_ready" + HOST_WAIT = "host_wait" + + +class NotificationCompletion(str, Enum): + """Meaning of terminal standalone-notification Work.""" + + LOCAL_ACCEPTED = "local_accepted" + REMOTE_DELIVERED = "remote_delivered" + + +@dataclass(frozen=True, slots=True) +class EndpointIdentity: + """Stable logical endpoint plus one process-lifetime incarnation.""" + + endpoint_id: str + incarnation: str + + def __post_init__(self) -> None: + _require_nonempty_str(self.endpoint_id, "endpoint_id") + _require_nonempty_str(self.incarnation, "incarnation") + + +@dataclass(frozen=True, slots=True) +class Capabilities: + read: bool = False + write: bool = False + attached_notifications: bool = False + standalone_notifications: bool = False + raw_spans: bool = False + strided_regions: bool = False + manual_progress: bool = False + background_progress: bool = False + cuda_completion_event: bool = False + persistent_request_slots: bool = False + fused_request_slot_start_poll: bool = False + fused_request_slot_poll: bool = False + cuda_ordering_modes: frozenset[CudaOrderingMode] = frozenset( + {CudaOrderingMode.CALLER_READY} + ) + standalone_notification_completion: NotificationCompletion | None = None + deferred_successful_request_slot_recycle: bool = False + # Appended to preserve positional construction compatibility for every + # pre-existing capability field. New code should use keyword arguments. + request_slot_notification_overrides: bool = False + fused_prevalidated_submit_poll: bool = False + fused_work_poll: bool = False + grouped_notification_receive: bool = False + synchronous_notification_send: bool = False + + def __post_init__(self) -> None: + values = { + "read": self.read, + "write": self.write, + "attached_notifications": self.attached_notifications, + "standalone_notifications": self.standalone_notifications, + "raw_spans": self.raw_spans, + "strided_regions": self.strided_regions, + "manual_progress": self.manual_progress, + "background_progress": self.background_progress, + "cuda_completion_event": self.cuda_completion_event, + "persistent_request_slots": self.persistent_request_slots, + "request_slot_notification_overrides": ( + self.request_slot_notification_overrides + ), + "fused_prevalidated_submit_poll": (self.fused_prevalidated_submit_poll), + "fused_work_poll": self.fused_work_poll, + "grouped_notification_receive": self.grouped_notification_receive, + "synchronous_notification_send": self.synchronous_notification_send, + "fused_request_slot_start_poll": self.fused_request_slot_start_poll, + "fused_request_slot_poll": self.fused_request_slot_poll, + "deferred_successful_request_slot_recycle": ( + self.deferred_successful_request_slot_recycle + ), + } + for name, value in values.items(): + if not isinstance(value, bool): + raise TypeError(f"capability {name} must be a bool") + try: + ordering_modes = frozenset(self.cuda_ordering_modes) + except TypeError as error: + raise TypeError("cuda_ordering_modes must be an iterable") from error + if not ordering_modes or not all( + isinstance(mode, CudaOrderingMode) for mode in ordering_modes + ): + raise TypeError("cuda_ordering_modes must contain CudaOrderingMode values") + if CudaOrderingMode.CALLER_READY not in ordering_modes: + raise ValueError("cuda_ordering_modes must include caller_ready") + object.__setattr__(self, "cuda_ordering_modes", ordering_modes) + if self.fused_request_slot_start_poll and not self.persistent_request_slots: + raise ValueError( + "fused_request_slot_start_poll requires persistent_request_slots" + ) + if ( + self.request_slot_notification_overrides + and not self.persistent_request_slots + ): + raise ValueError( + "request_slot_notification_overrides requires " + "persistent_request_slots" + ) + if self.request_slot_notification_overrides and not self.attached_notifications: + raise ValueError( + "request_slot_notification_overrides requires " "attached_notifications" + ) + if self.fused_request_slot_poll and not self.persistent_request_slots: + raise ValueError( + "fused_request_slot_poll requires persistent_request_slots" + ) + if ( + self.deferred_successful_request_slot_recycle + and not self.persistent_request_slots + ): + raise ValueError( + "deferred_successful_request_slot_recycle requires " + "persistent_request_slots" + ) + completion = self.standalone_notification_completion + if self.standalone_notifications and not isinstance( + completion, NotificationCompletion + ): + raise TypeError( + "notification-capable backends must declare standalone " + "notification completion" + ) + if not self.standalone_notifications and completion is not None: + raise ValueError( + "standalone notification completion requires notifications" + ) + if self.grouped_notification_receive and not self.notifications: + raise ValueError("grouped_notification_receive requires notifications") + if self.synchronous_notification_send and not self.standalone_notifications: + raise ValueError( + "synchronous_notification_send requires standalone_notifications" + ) + + @property + def notifications(self) -> bool: + """Whether either attached or standalone notifications are available.""" + + return self.attached_notifications or self.standalone_notifications + + +@dataclass(frozen=True, slots=True) +class EndpointInfo: + backend: str + endpoint_id: str + name: str + incarnation: str + capabilities: Capabilities + + def __post_init__(self) -> None: + _require_nonempty_str(self.backend, "backend") + _require_nonempty_str(self.endpoint_id, "endpoint_id") + _require_nonempty_str(self.name, "name") + _require_nonempty_str(self.incarnation, "incarnation") + if not isinstance(self.capabilities, Capabilities): + raise TypeError("capabilities must be Capabilities") + + @property + def identity(self) -> EndpointIdentity: + return EndpointIdentity(self.endpoint_id, self.incarnation) + + +@dataclass(frozen=True, slots=True) +class BackendRegion: + """Provider-facing region with a core-owned lifetime anchor. + + Providers consume ``backend_descriptor`` and must treat ``owner`` as opaque. + For local regions the descriptor is the provider registration handle; for + remote regions it is ``ImportedRegistration`` or ``ImportedRawMemory``. + In an indexed plan, ``count`` compactly contributes that many independently + selectable logical blocks, ordered by increasing stride displacement. + """ + + owner: object = field(repr=False) + offset: int + nbytes: int + device: str + memory_type: str + backend_descriptor: object = field(repr=False) + remote: bool = False + stride: int | None = None + count: int | None = None + + def __post_init__(self) -> None: + if not _is_int(self.offset) or self.offset < 0: + raise InvalidRegionError("offset must be a non-negative integer") + if not _is_int(self.nbytes) or self.nbytes <= 0: + raise InvalidRegionError("nbytes must be a positive integer") + _require_nonempty_str(self.device, "device", InvalidRegionError) + _require_nonempty_str(self.memory_type, "memory_type", InvalidRegionError) + if not isinstance(self.remote, bool): + raise InvalidRegionError("remote must be a bool") + if (self.stride is None) != (self.count is None): + raise InvalidRegionError("stride and count must be specified together") + if self.stride is not None: + if not _is_int(self.stride) or self.stride <= 0: + raise InvalidRegionError("stride must be a positive integer") + if ( + isinstance(self.count, bool) + or not isinstance(self.count, int) + or self.count <= 0 + ): + raise InvalidRegionError("count must be a positive integer") + if self.stride < self.nbytes: + raise InvalidRegionError("stride must be at least nbytes") + + +@dataclass(frozen=True, slots=True) +class BackendNotification: + MAX_PAYLOAD_BYTES: ClassVar[int] = 1024 * 1024 + + payload: bytes = field(repr=False) + source_endpoint_id: str | None = None + source_incarnation: str | None = None + + def __post_init__(self) -> None: + payload = _immutable_bytes(self.payload, "payload") + if len(payload) > self.MAX_PAYLOAD_BYTES: + raise ValueError("notification payload exceeds the wire size limit") + object.__setattr__(self, "payload", payload) + if self.source_endpoint_id is not None: + _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") + if self.source_incarnation is not None: + _require_nonempty_str(self.source_incarnation, "source_incarnation") + if (self.source_endpoint_id is None) != (self.source_incarnation is None): + raise ValueError( + "notification source_endpoint_id and source_incarnation must be " + "specified together" + ) + + @property + def source(self) -> EndpointIdentity | None: + if self.source_endpoint_id is None or self.source_incarnation is None: + return None + return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) + + +@dataclass(frozen=True, slots=True) +class BackendNotificationBatch: + """Provider-facing notification group from one native source batch. + + Providers may return multiple batches for one source. Payload order is + significant within and across that source's batches; no total order is + promised between different sources. + """ + + MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES + + payloads: tuple[bytes, ...] = field(repr=False) + total_nbytes: int = field(init=False) + source_endpoint_id: str | None = None + source_incarnation: str | None = None + + def __post_init__(self) -> None: + payloads, total_nbytes = _immutable_notification_payloads( + self.payloads, self.MAX_PAYLOAD_BYTES + ) + if not payloads: + raise ValueError("notification batch must contain at least one payload") + object.__setattr__(self, "payloads", payloads) + object.__setattr__(self, "total_nbytes", total_nbytes) + if self.source_endpoint_id is not None: + _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") + if self.source_incarnation is not None: + _require_nonempty_str(self.source_incarnation, "source_incarnation") + if (self.source_endpoint_id is None) != (self.source_incarnation is None): + raise ValueError( + "notification source_endpoint_id and source_incarnation must be " + "specified together" + ) + + @property + def source(self) -> EndpointIdentity | None: + if self.source_endpoint_id is None or self.source_incarnation is None: + return None + return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) + + +@dataclass(frozen=True, slots=True) +class Notification: + MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES + + payload: bytes = field(repr=False) + source_endpoint_id: str | None = None + source_incarnation: str | None = None + + def __post_init__(self) -> None: + payload = _immutable_bytes(self.payload, "payload") + if len(payload) > self.MAX_PAYLOAD_BYTES: + raise ValueError("notification payload exceeds the wire size limit") + object.__setattr__(self, "payload", payload) + if self.source_endpoint_id is not None: + _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") + if self.source_incarnation is not None: + _require_nonempty_str(self.source_incarnation, "source_incarnation") + if (self.source_endpoint_id is None) != (self.source_incarnation is None): + raise ValueError( + "notification source_endpoint_id and source_incarnation must be " + "specified together" + ) + + @property + def source(self) -> EndpointIdentity | None: + if self.source_endpoint_id is None or self.source_incarnation is None: + return None + return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) + + +@dataclass(frozen=True, slots=True) +class NotificationBatch: + """Public notification group preserving one provider source boundary.""" + + MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES + + payloads: tuple[bytes, ...] = field(repr=False) + source_endpoint_id: str | None = None + source_incarnation: str | None = None + + def __post_init__(self) -> None: + payloads, _ = _immutable_notification_payloads( + self.payloads, self.MAX_PAYLOAD_BYTES + ) + if not payloads: + raise ValueError("notification batch must contain at least one payload") + object.__setattr__(self, "payloads", payloads) + if self.source_endpoint_id is not None: + _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") + if self.source_incarnation is not None: + _require_nonempty_str(self.source_incarnation, "source_incarnation") + if (self.source_endpoint_id is None) != (self.source_incarnation is None): + raise ValueError( + "notification source_endpoint_id and source_incarnation must be " + "specified together" + ) + + @property + def source(self) -> EndpointIdentity | None: + if self.source_endpoint_id is None or self.source_incarnation is None: + return None + return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) + + +def _notification_batch_from_backend( + batch: BackendNotificationBatch, +) -> NotificationBatch: + """Adopt one exact, already-validated provider batch without rescanning it.""" + + if type(batch) is not BackendNotificationBatch: + raise TypeError("expected an exact BackendNotificationBatch") + result = object.__new__(NotificationBatch) + object.__setattr__(result, "payloads", batch.payloads) + object.__setattr__(result, "source_endpoint_id", batch.source_endpoint_id) + object.__setattr__(result, "source_incarnation", batch.source_incarnation) + return result + + +@dataclass(frozen=True, slots=True) +class RegistrationMetadata: + registration_id: str + name: str + nbytes: int + device: str + memory_type: str + backend_descriptor: bytes = field(repr=False) + + def __post_init__(self) -> None: + _require_nonempty_str( + self.registration_id, "registration_id", InvalidMetadataError + ) + _require_nonempty_str(self.name, "name", InvalidMetadataError) + if not _is_int(self.nbytes) or self.nbytes <= 0: + raise InvalidMetadataError("registration nbytes must be positive") + _require_nonempty_str(self.device, "device", InvalidMetadataError) + _require_nonempty_str(self.memory_type, "memory_type", InvalidMetadataError) + descriptor = _immutable_bytes( + self.backend_descriptor, "backend_descriptor", InvalidMetadataError + ) + object.__setattr__(self, "backend_descriptor", descriptor) + + +@dataclass(frozen=True, slots=True) +class EndpointMetadata: + WIRE_MAJOR: ClassVar[int] = 1 + WIRE_MINOR: ClassVar[int] = 0 + KNOWN_REQUIRED_FEATURES: ClassVar[frozenset[str]] = frozenset() + MAX_ENCODED_BYTES: ClassVar[int] = 64 * 1024 * 1024 + MAX_BACKEND_PAYLOAD_BYTES: ClassVar[int] = 32 * 1024 * 1024 + MAX_DESCRIPTOR_BYTES: ClassVar[int] = 4 * 1024 * 1024 + MAX_TOTAL_DESCRIPTOR_BYTES: ClassVar[int] = 32 * 1024 * 1024 + MAX_REGISTRATIONS: ClassVar[int] = 65_536 + MAX_REQUIRED_FEATURES: ClassVar[int] = 256 + MAX_EXTENSIONS: ClassVar[int] = 256 + MAX_EXTENSION_BYTES: ClassVar[int] = 4 * 1024 * 1024 + MAX_TOTAL_EXTENSION_BYTES: ClassVar[int] = 32 * 1024 * 1024 + MAX_STRING_BYTES: ClassVar[int] = 4_096 + + wire_major: int + wire_minor: int + required_features: tuple[str, ...] + backend: str + endpoint_id: str + name: str + incarnation: str + backend_payload: bytes = field(repr=False) + registrations: tuple[RegistrationMetadata, ...] + extensions: tuple[tuple[str, bytes], ...] = field(default=(), repr=False) + + def __post_init__(self) -> None: + if not _is_int(self.wire_major) or self.wire_major != self.WIRE_MAJOR: + raise InvalidMetadataError( + f"unsupported endpoint metadata wire major: {self.wire_major!r}" + ) + if not _is_int(self.wire_minor) or self.wire_minor < 0: + raise InvalidMetadataError( + f"invalid endpoint metadata wire minor: {self.wire_minor!r}" + ) + try: + required_features = tuple(self.required_features) + except TypeError as error: + raise InvalidMetadataError( + "required_features must be an iterable" + ) from error + if not all(isinstance(item, str) and item for item in required_features): + raise InvalidMetadataError( + "required_features must contain non-empty strings" + ) + if len(required_features) > self.MAX_REQUIRED_FEATURES: + raise InvalidMetadataError("too many required endpoint metadata features") + if any( + _utf8_nbytes(item, "required feature") > self.MAX_STRING_BYTES + for item in required_features + ): + raise InvalidMetadataError( + "required endpoint metadata feature exceeds the wire size limit" + ) + if len(required_features) != len(set(required_features)): + raise InvalidMetadataError("required_features must be unique") + unknown_features = set(required_features) - self.KNOWN_REQUIRED_FEATURES + if unknown_features: + raise InvalidMetadataError( + "unsupported required endpoint metadata features: " + f"{', '.join(sorted(unknown_features))}" + ) + object.__setattr__(self, "required_features", tuple(sorted(required_features))) + _require_nonempty_str(self.backend, "backend", InvalidMetadataError) + _require_nonempty_str(self.endpoint_id, "endpoint_id", InvalidMetadataError) + _require_nonempty_str(self.name, "name", InvalidMetadataError) + _require_nonempty_str(self.incarnation, "incarnation", InvalidMetadataError) + for field_name in ("backend", "endpoint_id", "name", "incarnation"): + if ( + _utf8_nbytes(getattr(self, field_name), field_name) + > self.MAX_STRING_BYTES + ): + raise InvalidMetadataError(f"{field_name} exceeds the wire size limit") + payload = _immutable_bytes( + self.backend_payload, "backend_payload", InvalidMetadataError + ) + if len(payload) > self.MAX_BACKEND_PAYLOAD_BYTES: + raise InvalidMetadataError("backend_payload exceeds the wire size limit") + object.__setattr__(self, "backend_payload", payload) + + try: + registrations = tuple(self.registrations) + except TypeError as error: + raise InvalidMetadataError("registrations must be an iterable") from error + if not all(isinstance(item, RegistrationMetadata) for item in registrations): + raise InvalidMetadataError( + "registrations must contain RegistrationMetadata values" + ) + registration_ids = [item.registration_id for item in registrations] + names = [item.name for item in registrations] + if len(registration_ids) != len(set(registration_ids)): + raise InvalidMetadataError("registration_id values must be unique") + if len(names) != len(set(names)): + raise InvalidMetadataError("registration names must be unique") + if len(registrations) > self.MAX_REGISTRATIONS: + raise InvalidMetadataError("too many endpoint metadata registrations") + descriptor_bytes = 0 + for registration in registrations: + for field_name in ( + "registration_id", + "name", + "device", + "memory_type", + ): + if ( + _utf8_nbytes( + getattr(registration, field_name), + f"registration {field_name}", + ) + > self.MAX_STRING_BYTES + ): + raise InvalidMetadataError( + f"registration {field_name} exceeds the wire size limit" + ) + if len(registration.backend_descriptor) > self.MAX_DESCRIPTOR_BYTES: + raise InvalidMetadataError( + "registration descriptor exceeds the wire size limit" + ) + descriptor_bytes += len(registration.backend_descriptor) + if descriptor_bytes > self.MAX_TOTAL_DESCRIPTOR_BYTES: + raise InvalidMetadataError( + "registration descriptors exceed the aggregate wire size limit" + ) + object.__setattr__(self, "registrations", registrations) + + raw_extensions: Any = self.extensions + if isinstance(raw_extensions, Mapping): + extension_items = tuple(raw_extensions.items()) + else: + try: + extension_items = tuple(raw_extensions) + except TypeError as error: + raise InvalidMetadataError("extensions must be a mapping") from error + if len(extension_items) > self.MAX_EXTENSIONS: + raise InvalidMetadataError("too many endpoint metadata extensions") + extensions: list[tuple[str, bytes]] = [] + extension_bytes = 0 + for item in extension_items: + if not isinstance(item, tuple) or len(item) != 2: + raise InvalidMetadataError("extensions must contain key/value pairs") + key, value = item + if not isinstance(key, str) or not key: + raise InvalidMetadataError("extension names must be non-empty strings") + if _utf8_nbytes(key, "extension name") > self.MAX_STRING_BYTES: + raise InvalidMetadataError("extension name exceeds the wire size limit") + extension = _immutable_bytes( + value, f"extension {key!r}", InvalidMetadataError + ) + if len(extension) > self.MAX_EXTENSION_BYTES: + raise InvalidMetadataError( + f"extension {key!r} exceeds the wire size limit" + ) + extension_bytes += len(extension) + extensions.append((key, extension)) + if extension_bytes > self.MAX_TOTAL_EXTENSION_BYTES: + raise InvalidMetadataError( + "endpoint metadata extensions exceed the aggregate wire size limit" + ) + extension_names = [key for key, _ in extensions] + if len(extension_names) != len(set(extension_names)): + raise InvalidMetadataError("extension names must be unique") + object.__setattr__(self, "extensions", tuple(sorted(extensions))) + + def __bytes__(self) -> bytes: + registrations = [ + { + "backend_descriptor": _encode_bytes(item.backend_descriptor), + "device": item.device, + "memory_type": item.memory_type, + "name": item.name, + "nbytes": item.nbytes, + "registration_id": item.registration_id, + } + for item in self.registrations + ] + value = { + "backend": self.backend, + "backend_payload": _encode_bytes(self.backend_payload), + "endpoint_id": self.endpoint_id, + "extensions": { + key: _encode_bytes(extension) for key, extension in self.extensions + }, + "incarnation": self.incarnation, + "name": self.name, + "registrations": registrations, + "required_features": list(self.required_features), + "wire_version": { + "major": self.wire_major, + "minor": self.wire_minor, + }, + } + encoded = json.dumps( + value, ensure_ascii=True, separators=(",", ":"), sort_keys=True + ).encode("utf-8") + if len(encoded) > self.MAX_ENCODED_BYTES: + raise InvalidMetadataError("endpoint metadata exceeds the wire size limit") + return encoded + + def to_bytes(self) -> bytes: + """Return the canonical wire representation.""" + + return bytes(self) + + @classmethod + def from_bytes(cls, data: bytes | bytearray | memoryview) -> "EndpointMetadata": + if not isinstance(data, (bytes, bytearray, memoryview)): + raise InvalidMetadataError("endpoint metadata must be bytes-like") + try: + encoded_nbytes = data.nbytes if isinstance(data, memoryview) else len(data) + except ValueError as error: + raise InvalidMetadataError( + "endpoint metadata buffer is unavailable" + ) from error + if encoded_nbytes > cls.MAX_ENCODED_BYTES: + raise InvalidMetadataError("endpoint metadata exceeds the wire size limit") + try: + encoded = bytes(data) + if len(encoded) > cls.MAX_ENCODED_BYTES: + raise InvalidMetadataError( + "endpoint metadata exceeds the wire size limit" + ) + value = json.loads(encoded.decode("utf-8"), object_pairs_hook=_json_object) + except ( + UnicodeDecodeError, + ValueError, + RecursionError, + InvalidMetadataError, + ) as error: + if isinstance(error, InvalidMetadataError): + raise + raise InvalidMetadataError("endpoint metadata is not valid JSON") from error + + _require_object_keys( + value, + { + "backend", + "backend_payload", + "endpoint_id", + "extensions", + "incarnation", + "name", + "registrations", + "required_features", + "wire_version", + }, + "endpoint metadata", + ) + wire_version = value["wire_version"] + _require_object_keys(wire_version, {"major", "minor"}, "wire_version") + required_features = value["required_features"] + if not isinstance(required_features, list): + raise InvalidMetadataError("required_features must be a JSON array") + if len(required_features) > cls.MAX_REQUIRED_FEATURES: + raise InvalidMetadataError("too many required endpoint metadata features") + raw_extensions = value["extensions"] + if not isinstance(raw_extensions, dict): + raise InvalidMetadataError("extensions must be a JSON object") + if len(raw_extensions) > cls.MAX_EXTENSIONS: + raise InvalidMetadataError("too many endpoint metadata extensions") + raw_registrations = value["registrations"] + if not isinstance(raw_registrations, list): + raise InvalidMetadataError("registrations must be a JSON array") + if len(raw_registrations) > cls.MAX_REGISTRATIONS: + raise InvalidMetadataError("too many endpoint metadata registrations") + registrations = tuple( + _registration_from_json(item, index) + for index, item in enumerate(raw_registrations) + ) + return cls( + wire_major=wire_version["major"], + wire_minor=wire_version["minor"], + required_features=tuple(required_features), + backend=value["backend"], + endpoint_id=value["endpoint_id"], + name=value["name"], + incarnation=value["incarnation"], + backend_payload=_decode_bytes(value["backend_payload"], "backend_payload"), + registrations=registrations, + extensions=tuple( + (key, _decode_bytes(extension, f"extension {key!r}")) + for key, extension in raw_extensions.items() + ), + ) + + +def _registration_from_json(value: Any, index: int) -> RegistrationMetadata: + context = f"registration at index {index}" + _require_object_keys( + value, + { + "backend_descriptor", + "device", + "memory_type", + "name", + "nbytes", + "registration_id", + }, + context, + ) + return RegistrationMetadata( + registration_id=value["registration_id"], + name=value["name"], + nbytes=value["nbytes"], + device=value["device"], + memory_type=value["memory_type"], + backend_descriptor=_decode_bytes( + value["backend_descriptor"], f"{context} backend_descriptor" + ), + ) + + +def _json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise InvalidMetadataError(f"duplicate JSON field: {key!r}") + value[key] = item + return value + + +def _require_object_keys(value: Any, expected: set[str], context: str) -> None: + if not isinstance(value, dict): + raise InvalidMetadataError(f"{context} must be a JSON object") + keys = set(value) + missing = expected - keys + unknown = keys - expected + if missing: + raise InvalidMetadataError( + f"{context} is missing fields: {', '.join(sorted(missing))}" + ) + if unknown: + raise InvalidMetadataError( + f"{context} has unknown fields: {', '.join(sorted(unknown))}" + ) + + +def _encode_bytes(value: bytes) -> str: + return base64.b64encode(value).decode("ascii") + + +def _decode_bytes(value: Any, field: str) -> bytes: + if not isinstance(value, str): + raise InvalidMetadataError(f"{field} must be a base64 string") + try: + encoded = value.encode("ascii") + return base64.b64decode(encoded, validate=True) + except (UnicodeEncodeError, binascii.Error, ValueError) as error: + raise InvalidMetadataError(f"{field} is not valid base64") from error + + +def _immutable_bytes( + value: Any, + field: str, + error_type: type[Exception] = TypeError, +) -> bytes: + if not isinstance(value, (bytes, bytearray, memoryview)): + raise error_type(f"{field} must be bytes-like") + return bytes(value) + + +def _immutable_notification_payloads( + value: Any, + max_payload_bytes: int, +) -> tuple[tuple[bytes, ...], int]: + """Snapshot payloads and count bytes during the required size validation.""" + + try: + payloads = value if type(value) is tuple else tuple(value) + except TypeError as error: + raise TypeError("payloads must be an iterable of bytes-like values") from error + total_nbytes = 0 + for payload in payloads: + if type(payload) is not bytes: + for item in payloads: + if not isinstance(item, (bytes, bytearray, memoryview)): + raise TypeError("payload must be bytes-like") + reported_nbytes = ( + item.nbytes if isinstance(item, memoryview) else len(item) + ) + if reported_nbytes > max_payload_bytes: + raise ValueError("notification payload exceeds the wire size limit") + immutable_payloads = tuple( + _immutable_bytes(item, "payload") for item in payloads + ) + total_nbytes = 0 + for item in immutable_payloads: + item_nbytes = len(item) + if item_nbytes > max_payload_bytes: + raise ValueError("notification payload exceeds the wire size limit") + total_nbytes += item_nbytes + return immutable_payloads, total_nbytes + payload_nbytes = len(payload) + if payload_nbytes > max_payload_bytes: + raise ValueError("notification payload exceeds the wire size limit") + total_nbytes += payload_nbytes + return payloads, total_nbytes + + +def _require_nonempty_str( + value: Any, + field: str, + error_type: type[Exception] = TypeError, +) -> None: + if not isinstance(value, str) or not value: + raise error_type(f"{field} must be a non-empty string") + + +def _utf8_nbytes(value: str, field: str) -> int: + try: + return len(value.encode("utf-8")) + except UnicodeEncodeError as error: + raise InvalidMetadataError(f"{field} is not valid Unicode") from error + + +def _is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) From 5c8c0489ea02aa3ea5cadd73e5dd077fb4fb2135 Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Thu, 17 Sep 2026 16:12:09 +0200 Subject: [PATCH 2/2] distributed: reduce transfer prototype to review MVP Keep ownership, prepared catalogs, READ/WRITE and explicit readiness. Remove compiled dispatch, recovery machinery and unrelated changes. Success-path prototype only; native validation remains pending. Assisted-by: OpenAI Codex Signed-off-by: Samuel Nordmann --- build_variables.bzl | 1 - test/distributed/test_transfer.py | 9292 +---------------- test/dynamo/test_trace_rules.py | 1 - torch/_C/_distributed_transfer.pyi | 7 - torch/csrc/Module.cpp | 2 - .../distributed/transfer/python_bindings.cpp | 432 - .../distributed/transfer/python_bindings.h | 9 - torch/distributed/_transfer/README.md | 856 +- torch/distributed/_transfer/__init__.py | 124 +- torch/distributed/_transfer/_api.py | 8432 +-------------- torch/distributed/_transfer/_backend.py | 1057 +- torch/distributed/_transfer/_errors.py | 62 - torch/distributed/_transfer/_reference.py | 814 +- torch/distributed/_transfer/_types.py | 867 -- 14 files changed, 482 insertions(+), 21474 deletions(-) delete mode 100644 torch/_C/_distributed_transfer.pyi delete mode 100644 torch/csrc/distributed/transfer/python_bindings.cpp delete mode 100644 torch/csrc/distributed/transfer/python_bindings.h delete mode 100644 torch/distributed/_transfer/_errors.py delete mode 100644 torch/distributed/_transfer/_types.py diff --git a/build_variables.bzl b/build_variables.bzl index 10b166d65fd20..9591da3dbed6d 100644 --- a/build_variables.bzl +++ b/build_variables.bzl @@ -975,7 +975,6 @@ libtorch_python_core_sources = [ "torch/csrc/autograd/python_variable_indexing.cpp", "torch/csrc/PyObjectConversionPythonImpl.cpp", "torch/csrc/distributed/python_placement.cpp", - "torch/csrc/distributed/transfer/python_bindings.cpp", "torch/csrc/dynamo/python_compiled_autograd.cpp", "torch/csrc/dynamo/cache_entry.cpp", "torch/csrc/dynamo/cpp_shim.cpp", diff --git a/test/distributed/test_transfer.py b/test/distributed/test_transfer.py index b1aaa457eb93f..1d75962e9acdc 100644 --- a/test/distributed/test_transfer.py +++ b/test/distributed/test_transfer.py @@ -1,9211 +1,127 @@ -import ast import gc -from array import array -from pathlib import Path -import sys -import threading +import unittest import weakref -from unittest import mock +from unittest.mock import Mock import torch -from torch.distributed._transfer import ( - BACKEND_FACTORY_API_VERSION, - BackendFailureError, - BackendIndexSelection, - BackendNotificationBatch, - BackendRequestSlot, - BackendRegion, - BackendWork, - BoundTransfer, - BusyError, - Capabilities, - ClosedError, - CudaOrderingMode, - Endpoint, - EndpointIdentity, - EndpointMetadata, - ImportedRawMemory, - ImportedRegistration, - IndexSelection, - InvalidMetadataError, - InvalidRegionError, - NotificationCompletion, - NotificationBatch, - RawSpan, - RequestLocalNotificationSlot, - RegisteredMemory, - RegistrationMetadata, - TransferCancelledError, - TransferOp, - TransferSlot, - TransferTimeoutError, - ThreadMode, - UnsupportedError, - Work, - WorkState, - register_backend, -) -from torch.distributed._transfer import _api, _backend, _reference, _types -from torch.distributed._transfer._reference import _ReferenceBackend -from torch.testing._internal.common_utils import run_tests, TestCase +from torch.distributed._transfer import Endpoint +from torch.distributed._transfer._reference import MemoryBackend -def _emulated_compiled_index_analyzer( - local_indices: object, - remote_indices: object, - destination_is_remote: bool, -) -> tuple[bytes, int, int, int, int, int, int, bool]: - local_bytes = bytes(memoryview(local_indices).cast("B")) - remote_bytes = bytes(memoryview(remote_indices).cast("B")) - packed = local_bytes + remote_bytes - values = memoryview(packed).cast("i") - local_count = len(local_bytes) // 4 - remote_count = len(remote_bytes) // 4 - local = values[:local_count] - remote = values[local_count:] - destination = remote if destination_is_remote else local - return ( - packed, - local_count, - remote_count, - min(local), - max(local), - min(remote), - max(remote), - len(set(destination)) != len(destination), - ) - - -class _PendingWork(BackendWork): - def __init__(self) -> None: - self._state = WorkState.PENDING - - @property - def state(self) -> WorkState: - return self._state - - @property - def error(self) -> BaseException | None: - return None - - def cancel(self) -> bool: - self._state = WorkState.CANCELLED - return True - - def close(self) -> None: - pass - - -class _BadErrorWork(_PendingWork): - @property - def state(self) -> WorkState: - return WorkState.FAILED - - @property - def error(self): - return "not an exception" - - -class _RegressingWork(_PendingWork): - def __init__(self) -> None: - self._queries = 0 - self._cancelled = False - - @property - def state(self) -> WorkState: - if self._cancelled: - return WorkState.CANCELLED - self._queries += 1 - return WorkState.PENDING if self._queries == 2 else WorkState.RUNNING - - def cancel(self) -> bool: - self._cancelled = True - return True - - -class _TransientStateErrorWork(_PendingWork): - def __init__(self) -> None: - self._queries = 0 - self._cancelled = False - - @property - def state(self) -> WorkState: - self._queries += 1 - if self._queries == 1: - raise RuntimeError("transient query failure") - if self._cancelled: - return WorkState.CANCELLED - return WorkState.RUNNING - - def cancel(self) -> bool: - self._cancelled = True - return True - - -class _CountingPendingWork(_PendingWork): - def __init__(self) -> None: - super().__init__() - self.state_queries = 0 - - @property - def state(self) -> WorkState: - self.state_queries += 1 - return self._state - - -class _InterruptingFailureWork(_PendingWork): - def __init__(self) -> None: - super().__init__() - self.failure = RuntimeError("provider failure") - self.error_queries = 0 - - @property - def state(self) -> WorkState: - return WorkState.FAILED - - @property - def error(self) -> BaseException: - self.error_queries += 1 - if self.error_queries == 1: - raise KeyboardInterrupt - return self.failure - - -class _PendingReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - kwargs["adopt_work"](_PendingWork()) - - -class _CountingPendingReferenceBackend(_ReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.last_work: _CountingPendingWork | None = None - - def submit(self, *args, **kwargs) -> None: - self.last_work = _CountingPendingWork() - kwargs["adopt_work"](self.last_work) - - -class _InterruptingFailureReferenceBackend(_ReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.last_work: _InterruptingFailureWork | None = None - - def submit(self, *args, **kwargs) -> None: - self.last_work = _InterruptingFailureWork() - kwargs["adopt_work"](self.last_work) - - -class _FusedSelectionWork(_PendingWork): - def __init__(self, state: WorkState) -> None: - super().__init__() - self._state = state - self.state_queries = 0 - self.bounded_calls: list[tuple[int, int | None, int]] = [] - self.next_bounded_state = WorkState.COMPLETED - self._error: BaseException | None = None - - @property - def state(self) -> WorkState: - self.state_queries += 1 - return self._state - - @property - def error(self) -> BaseException | None: - return self._error - - def poll_bounded( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState: - self.bounded_calls.append((max_polls, timeout_ns, timeout_check_interval)) - self._state = self.next_bounded_state - return self._state - - -class _FusedSelectionReferenceBackend(_ReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - fused_prevalidated_submit_poll=True, - fused_work_poll=True, - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.initial_state = WorkState.RUNNING - self.fused_calls: list[tuple[int, int | None, bytes | None]] = [] - self.last_work: _FusedSelectionWork | None = None - - def submit_prevalidated_and_poll( - self, - op, - plan_handle, - *, - selection, - notification, - max_polls, - timeout_ns, - adopt_work, - ) -> WorkState: - selection._validate_binding(op, plan_handle) - work = _FusedSelectionWork(self.initial_state) - self.last_work = work - adopt_work(work) - self.fused_calls.append((max_polls, timeout_ns, notification)) - return self.initial_state - - -class _BackgroundOnlyReferenceBackend(_ReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - background_progress=True, - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - def progress(self) -> None: - raise AssertionError("background-only backend must not be progressed") - - -class _AttachedOnlyReferenceBackend(_ReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - ) - - -class _BadErrorReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - kwargs["adopt_work"](_BadErrorWork()) - - -class _InterruptReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - raise KeyboardInterrupt - - -class _AdoptThenInterruptReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - kwargs["adopt_work"](_PendingWork()) - raise KeyboardInterrupt - - -class _RegressingReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - kwargs["adopt_work"](_RegressingWork()) - - -class _TransientStateErrorReferenceBackend(_ReferenceBackend): - def submit(self, *args, **kwargs) -> None: - kwargs["adopt_work"](_TransientStateErrorWork()) - - -class _InterruptAfterCloseReferenceBackend(_ReferenceBackend): - def close(self) -> None: - was_closed = self._closed - super().close() - if not was_closed: - raise KeyboardInterrupt - - -class _TestRequestSlot(BackendRequestSlot): - def __init__(self) -> None: - self._state: WorkState | None = None - self._error: BaseException | None = None - self.next_state = WorkState.COMPLETED - self.next_error: BaseException | None = None - self.start_count = 0 - self.recycle_count = 0 - self.close_count = 0 - self.closed = False - self.fail_close_once = False - self.raise_error_once: BaseException | None = None - self.start_entered: threading.Event | None = None - self.start_release: threading.Event | None = None - self.completion = object() - - @property - def state(self) -> WorkState | None: - if self.closed: - raise ClosedError("test request slot is closed") - return self._state - - @property - def error(self) -> BaseException | None: - if self.closed: - raise ClosedError("test request slot is closed") - if self.raise_error_once is not None: - error = self.raise_error_once - self.raise_error_once = None - raise error - return self._error - - @property - def completion_event(self) -> object: - return self.completion - - def start(self) -> None: - if self.closed: - raise ClosedError("test request slot is closed") - self.start_count += 1 - if self.start_entered is not None: - self.start_entered.set() - if self.start_release is not None and not self.start_release.wait(2): - raise RuntimeError("test request slot start timed out") - self._state = self.next_state - self._error = self.next_error - - def cancel(self) -> bool: - if self.closed: - raise ClosedError("test request slot is closed") - if self._state in { - WorkState.COMPLETED, - WorkState.FAILED, - WorkState.CANCELLED, - }: - return self._state is WorkState.CANCELLED - self._state = WorkState.CANCELLED - self._error = None - return True - - def recycle(self) -> None: - if self.closed: - raise ClosedError("test request slot is closed") - self.recycle_count += 1 - self._state = None - self._error = None - - def close(self) -> None: - self.close_count += 1 - if self.fail_close_once: - self.fail_close_once = False - raise RuntimeError("injected request-slot cleanup failure") - self.closed = True - - -class _UnacknowledgedStartRequestSlot(_TestRequestSlot): - @property - def state(self) -> WorkState | None: - return None - - def start(self) -> None: - raise KeyboardInterrupt - - -class _NormallyUnacknowledgedRequestSlot(_UnacknowledgedStartRequestSlot): - def start(self) -> None: - self.start_count += 1 - - -class _BlockingObservationRequestSlot(_TestRequestSlot): - def __init__(self) -> None: - super().__init__() - self.block_next_observation = False - self.observation_entered = threading.Event() - self.observation_release = threading.Event() - - @property - def state(self) -> WorkState | None: - if self.block_next_observation: - self.block_next_observation = False - self.observation_entered.set() - if not self.observation_release.wait(2): - raise RuntimeError("test request-slot observation timed out") - return super().state - - -_RETURN_CURRENT_SLOT_STATE = object() - - -class _FusedTestRequestSlot(_TestRequestSlot): - def __init__(self) -> None: - super().__init__() - self.fused_calls: list[tuple[int, int | None]] = [] - self.fused_result: object = _RETURN_CURRENT_SLOT_STATE - self.fused_error: BaseException | None = None - self.poll_states: list[WorkState | None] = [] - self.poll_count = 0 - self.bounded_calls: list[tuple[int, int | None, int]] = [] - self.bounded_results: list[object] = [] - self.bounded_error: BaseException | None = None - self.bounded_entered: threading.Event | None = None - self.bounded_release: threading.Event | None = None - - def poll_state(self) -> WorkState | None: - self.poll_count += 1 - if self.poll_states: - state = self.poll_states.pop(0) - if isinstance(state, WorkState): - self._state = state - return state - return self.state - - def start_and_poll( - self, - *, - max_polls: int, - timeout_ns: int | None, - ): - self.fused_calls.append((max_polls, timeout_ns)) - self.start() - if self.fused_error is not None: - raise self.fused_error - if self.fused_result is _RETURN_CURRENT_SLOT_STATE: - return self._state - return self.fused_result - - def poll_bounded( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ): - self.bounded_calls.append((max_polls, timeout_ns, timeout_check_interval)) - if self.bounded_entered is not None: - self.bounded_entered.set() - if self.bounded_release is not None and not self.bounded_release.wait(2): - raise RuntimeError("test bounded request-slot poll timed out") - if self.bounded_error is not None: - raise self.bounded_error - if self.bounded_results: - state = self.bounded_results.pop(0) - if isinstance(state, WorkState): - self._state = state - return state - return super().poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - - -class _NotificationFusedTestRequestSlot(_FusedTestRequestSlot): - def __init__(self) -> None: - super().__init__() - self.notification_starts: list[bytes | None] = [] - self.notification_fused_calls: list[tuple[bytes | None, int, int | None]] = [] - - def start_with_notification(self, notification: bytes | None) -> None: - self.notification_starts.append(notification) - self.start() - - def start_and_poll_with_notification( - self, - notification: bytes | None, - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - self.notification_fused_calls.append((notification, max_polls, timeout_ns)) - self.start() - if self.fused_error is not None: - raise self.fused_error - if self.fused_result is _RETURN_CURRENT_SLOT_STATE: - return self._state - return self.fused_result - - -class _TestExecutionDispatch: - _OPERATION_NAMES = frozenset( - { - "start", - "start_and_poll", - "poll_bounded", - "poll_state", - "cancel", - "recycle", - } - ) - - def __init__(self, slot: _FusedTestRequestSlot) -> None: - self.slot = slot - self.attribute_reads = {name: 0 for name in self._OPERATION_NAMES} - - def __getattribute__(self, name): - if name in object.__getattribute__(self, "_OPERATION_NAMES"): - reads = object.__getattribute__(self, "attribute_reads") - reads[name] += 1 - return object.__getattribute__(self, name) - - def start(self) -> None: - self.slot.start() - - def start_and_poll(self, *, max_polls: int, timeout_ns: int | None): - return self.slot.start_and_poll( - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - - def poll_bounded( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ): - return self.slot.poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - - def poll_state(self) -> WorkState | None: - return self.slot.poll_state() - - def cancel(self) -> bool: - return self.slot.cancel() - - def recycle(self) -> None: - self.slot.recycle() - - @property - def error(self): - raise AssertionError("execution dispatch must not own diagnostics") - - @property - def completion_event(self): - raise AssertionError("execution dispatch must not own completion events") - - def close(self) -> None: - raise AssertionError("execution dispatch must not own close") - - -class _CapabilityGatedExecutionDispatch(_TestExecutionDispatch): - def __getattribute__(self, name): - if name in {"start_and_poll", "poll_bounded"}: - raise AssertionError("disabled fused dispatch operation was resolved") - return super().__getattribute__(name) - - -class _NonCallableExecutionDispatch(_TestExecutionDispatch): - start = 1 - - -class _NoneStartExecutionDispatch(_TestExecutionDispatch): - start = None - - -class _NoneFusedExecutionDispatch(_TestExecutionDispatch): - start_and_poll = None - - -class _UnretainedExecutionDispatch: - def __init__(self, slot: _FusedTestRequestSlot) -> None: - # None of these callable objects retains this dispatch container. - self.start = slot.start - self.start_and_poll = slot.start_and_poll - self.poll_bounded = slot.poll_bounded - self.poll_state = slot.poll_state - self.cancel = slot.cancel - self.recycle = slot.recycle - - -class _UnacknowledgedFusedRequestSlot(_FusedTestRequestSlot): - @property - def state(self) -> WorkState | None: - return None - - def poll_state(self) -> WorkState | None: - self.poll_count += 1 - return None - - def start_and_poll( - self, - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - self.fused_calls.append((max_polls, timeout_ns)) - self.start_count += 1 - raise KeyboardInterrupt - - -class _SlotReferenceBackend(_ReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.request_slots: list[_TestRequestSlot] = [] - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _TestRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.request_slots.append(slot) - - -class _FusedSlotReferenceBackend(_SlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - fused_request_slot_start_poll=True, - fused_request_slot_poll=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.request_slots.append(slot) - - -class _BlockingObservationSlotReferenceBackend(_SlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _BlockingObservationRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.request_slots.append(slot) - - -class _ExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.execution_dispatches: list[_TestExecutionDispatch] = [] - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - dispatch = _TestExecutionDispatch(slot) - try: - kwargs["adopt_slot"](slot, execution_dispatch=dispatch) - except BaseException: - return - self.request_slots.append(slot) - self.execution_dispatches.append(dispatch) - - -class _CapabilityGatedDispatchSlotBackend(_SlotReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.execution_dispatches: list[_CapabilityGatedExecutionDispatch] = [] - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - dispatch = _CapabilityGatedExecutionDispatch(slot) - try: - kwargs["adopt_slot"](slot, execution_dispatch=dispatch) - except BaseException: - return - self.request_slots.append(slot) - self.execution_dispatches.append(dispatch) - - -class _InvalidExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): - _DISPATCH_TYPE = _NonCallableExecutionDispatch - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - dispatch = self._DISPATCH_TYPE(slot) - try: - kwargs["adopt_slot"](slot, execution_dispatch=dispatch) - except BaseException: - return - raise AssertionError("non-callable execution dispatch was adopted") - - -class _NoneStartExecutionDispatchSlotBackend(_InvalidExecutionDispatchSlotBackend): - _DISPATCH_TYPE = _NoneStartExecutionDispatch - - -class _NoneFusedExecutionDispatchSlotBackend(_InvalidExecutionDispatchSlotBackend): - _DISPATCH_TYPE = _NoneFusedExecutionDispatch - - -class _UnretainedExecutionDispatchSlotBackend(_FusedSlotReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.execution_dispatch_ref = None - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - dispatch = _UnretainedExecutionDispatch(slot) - self.execution_dispatch_ref = weakref.ref(dispatch) - try: - kwargs["adopt_slot"](slot, execution_dispatch=dispatch) - except BaseException: - return - self.request_slots.append(slot) - - -class _StartFusedOnlySlotReferenceBackend(_FusedSlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - fused_request_slot_start_poll=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - -class _PollOnlySlotReferenceBackend(_FusedSlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - fused_request_slot_poll=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - -class _DeferredRecycleSlotReferenceBackend(_SlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - deferred_successful_request_slot_recycle=True, - ) - - -class _DeferredRecycleFusedSlotReferenceBackend(_FusedSlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - fused_request_slot_start_poll=True, - fused_request_slot_poll=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - deferred_successful_request_slot_recycle=True, - ) - - -class _NotificationSlotReferenceBackend(_FusedSlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - fused_request_slot_start_poll=True, - fused_request_slot_poll=True, - cuda_completion_event=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - deferred_successful_request_slot_recycle=True, - request_slot_notification_overrides=True, - ) - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _NotificationFusedTestRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.request_slots.append(slot) - - -class _MissingNotificationSlotBackend(_NotificationSlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _FusedTestRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - raise AssertionError("default notification dispatch was adopted") - - -class _UnacknowledgedFusedSlotBackend(_FusedSlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _UnacknowledgedFusedRequestSlot() - kwargs["adopt_slot"](slot) - self.request_slots.append(slot) - - -class _UnacknowledgedStartSlotBackend(_SlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _UnacknowledgedStartRequestSlot() - kwargs["adopt_slot"](slot) - self.request_slots.append(slot) - - -class _NormallyUnacknowledgedSlotBackend(_SlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _NormallyUnacknowledgedRequestSlot() - kwargs["adopt_slot"](slot) - self.request_slots.append(slot) - - -class _InterruptSlotCreationBackend(_SlotReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.created_slot: _TestRequestSlot | None = None - - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _TestRequestSlot() - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.created_slot = slot - raise KeyboardInterrupt - - -class _InterruptSlotCleanupBackend(_InterruptSlotCreationBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args - slot = _TestRequestSlot() - slot.fail_close_once = True - try: - kwargs["adopt_slot"](slot) - except BaseException: - return - self.created_slot = slot - raise KeyboardInterrupt - - -class _InterruptBeforeSlotAdoptionBackend(_SlotReferenceBackend): - def create_request_slot(self, *args, **kwargs) -> None: - del args, kwargs - raise KeyboardInterrupt - - -class _NoCompletionEventSlotBackend(_SlotReferenceBackend): - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - persistent_request_slots=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - -register_backend("test-pending", _PendingReferenceBackend, replace=True) -register_backend( - "test-counting-pending", _CountingPendingReferenceBackend, replace=True -) -register_backend( - "test-interrupting-failure", - _InterruptingFailureReferenceBackend, - replace=True, -) -register_backend("test-fused-selection", _FusedSelectionReferenceBackend, replace=True) -register_backend("test-background-only", _BackgroundOnlyReferenceBackend, replace=True) -register_backend("test-attached-only", _AttachedOnlyReferenceBackend, replace=True) -register_backend("test-bad-error", _BadErrorReferenceBackend, replace=True) -register_backend("test-interrupt", _InterruptReferenceBackend, replace=True) -register_backend( - "test-adopt-then-interrupt", _AdoptThenInterruptReferenceBackend, replace=True -) -register_backend( - "test-request-slot-no-completion-event", - _NoCompletionEventSlotBackend, - replace=True, -) -register_backend("test-regressing", _RegressingReferenceBackend, replace=True) -register_backend( - "test-transient-state-error", - _TransientStateErrorReferenceBackend, - replace=True, -) -register_backend( - "test-interrupt-after-close", _InterruptAfterCloseReferenceBackend, replace=True -) -register_backend("test-request-slot", _SlotReferenceBackend, replace=True) -register_backend( - "test-request-slot-blocking-observation", - _BlockingObservationSlotReferenceBackend, - replace=True, -) -register_backend("test-request-slot-fused", _FusedSlotReferenceBackend, replace=True) -register_backend( - "test-request-slot-execution-dispatch", - _ExecutionDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-execution-dispatch-gated", - _CapabilityGatedDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-execution-dispatch-invalid", - _InvalidExecutionDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-execution-dispatch-none-start", - _NoneStartExecutionDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-execution-dispatch-none-fused", - _NoneFusedExecutionDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-execution-dispatch-lifetime", - _UnretainedExecutionDispatchSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-start-fused-only", - _StartFusedOnlySlotReferenceBackend, - replace=True, -) -register_backend( - "test-request-slot-poll-only", - _PollOnlySlotReferenceBackend, - replace=True, -) -register_backend( - "test-request-slot-deferred-recycle", - _DeferredRecycleSlotReferenceBackend, - replace=True, -) -register_backend( - "test-request-slot-fused-deferred-recycle", - _DeferredRecycleFusedSlotReferenceBackend, - replace=True, -) -register_backend( - "test-request-slot-notification", - _NotificationSlotReferenceBackend, - replace=True, -) -register_backend( - "test-request-slot-notification-missing", - _MissingNotificationSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-fused-unacknowledged", - _UnacknowledgedFusedSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-unacknowledged-start", - _UnacknowledgedStartSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-no-ack", - _NormallyUnacknowledgedSlotBackend, - replace=True, -) -register_backend( - "test-request-slot-create-interrupt", - _InterruptSlotCreationBackend, - replace=True, -) -register_backend( - "test-request-slot-cleanup-interrupt", - _InterruptSlotCleanupBackend, - replace=True, -) -register_backend( - "test-request-slot-before-adoption-interrupt", - _InterruptBeforeSlotAdoptionBackend, - replace=True, -) - - -class EndpointTransferTest(TestCase): - def test_backend_registration_lost_return_retry_is_exactly_idempotent( - self, - ) -> None: - class Factory: - def __call__(self, *args, adopt_backend, **kwargs) -> None: - del args, adopt_backend, kwargs - - backend_name = "test-registration-lost-return" - factory = Factory() - register_backend(backend_name, factory, replace=True, factory_api_version=2) - - # Model a retry after Core committed the registry write but the caller - # was interrupted before observing the return. - register_backend(backend_name, factory, factory_api_version=2) - registration = _backend._get_backend_registration(backend_name) - self.assertIs(registration.factory, factory) - self.assertEqual(registration.factory_api_version, 2) - - # Idempotence is deliberately exact; neither a new callable nor a - # different contract can adopt an already-published name implicitly. - with self.assertRaises(BusyError): - register_backend(backend_name, factory, factory_api_version=1) - with self.assertRaises(BusyError): - register_backend(backend_name, Factory(), factory_api_version=2) - - def test_backend_factory_versions_are_explicit_and_never_introspected(self) -> None: - self.assertEqual(BACKEND_FACTORY_API_VERSION, 2) - for value, error_type in ( - (True, TypeError), - (2.0, TypeError), - (0, ValueError), - (3, ValueError), - ): - with self.subTest(value=value), self.assertRaises(error_type): - register_backend( - f"test-invalid-factory-version-{value!r}", - lambda *args: None, - factory_api_version=value, - ) - - events = [] - - class Factory: - @property - def __signature__(self): - raise AssertionError("factory signatures must not be inspected") - - def __call__(self, *args, adopt_backend, **kwargs) -> None: - del kwargs - backend = _ReferenceBackend._new_unactivated(*args) - events.append("pure") - self.assert_not_activated(backend) - adopt_backend(backend) - events.append("adopted") - backend._activate() - events.append("activated") - - @staticmethod - def assert_not_activated(backend) -> None: - if backend._activated: - raise AssertionError("native/global activation preceded adoption") - - backend_name = "test-factory-v2-no-introspection" - register_backend(backend_name, Factory(), replace=True, factory_api_version=2) - endpoint = Endpoint( - "factory-v2", backend=backend_name, endpoint_id="factory-v2-order" - ) - self.assertEqual(events, ["pure", "adopted", "activated"]) - self.assertTrue(endpoint._backend_receipt.published) - endpoint.close() - - # The legacy callable contract remains the default. - legacy_name = "test-factory-v1-compatible" - register_backend(legacy_name, _ReferenceBackend, replace=True) - legacy = Endpoint("factory-v1", backend=legacy_name) - self.assertEqual( - _backend._get_backend_registration(legacy_name).factory_api_version, 1 - ) - legacy.close() - - def test_backend_factory_v2_rejects_invalid_adoption_contracts(self) -> None: - retained_callbacks = [] - - for behavior in ("missing", "duplicate", "wrong", "non-none"): - instances = [] - - def factory(*args, adopt_backend, _behavior=behavior, **kwargs): - del kwargs - if _behavior == "missing": - retained_callbacks.append(adopt_backend) - return None - if _behavior == "wrong": - adopt_backend(object()) - return None - backend = _ReferenceBackend._new_unactivated(*args) - instances.append(backend) - adopt_backend(backend) - if _behavior == "duplicate": - adopt_backend(backend) - backend._activate() - return object() if _behavior == "non-none" else None - - backend_name = f"test-factory-v2-invalid-{behavior}" - register_backend(backend_name, factory, replace=True, factory_api_version=2) - with ( - self.subTest(behavior=behavior), - self.assertRaises((BackendFailureError, TypeError)), - ): - Endpoint("invalid-v2", backend=backend_name) - for instance in instances: - self.assertTrue(instance._closed) - - late_backend = _ReferenceBackend( - "late-adoption", - "late-adoption", - "late-incarnation", - _types.ProgressMode.MANUAL, - ThreadMode.SERIALIZED, - {}, - ) - with self.assertRaisesRegex(BackendFailureError, "after its call returned"): - retained_callbacks[0](late_backend) - late_backend.close() - - def test_backend_factory_v1_wrong_type_is_foreign_and_never_touched(self) -> None: - calls = [] - - class ForeignResult: - def close(self) -> None: - calls.append("close") - - candidate = ForeignResult() - backend_name = "test-factory-v1-foreign-result" - register_backend( - backend_name, - lambda *args, **kwargs: candidate, - replace=True, - factory_api_version=1, - ) - - with self.assertRaisesRegex(TypeError, "must return a Backend instance"): - Endpoint("foreign-result", backend=backend_name) - self.assertEqual(calls, []) - - def test_backend_factory_v2_post_adoption_failure_exposes_recovery(self) -> None: - instances = [] - - class PendingCloseBackend(_ReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - - def close(self) -> None: - if not hasattr(self, "allow_close"): - self.allow_close = False - if not self.allow_close: - raise BusyError("injected activation cleanup hold") - super().close() - - def factory(*args, adopt_backend, **kwargs) -> None: - del kwargs - backend = PendingCloseBackend._new_unactivated(*args) - instances.append(backend) - adopt_backend(backend) - backend._activate() - raise KeyboardInterrupt("post-adoption activation cut") - - backend_name = "test-factory-v2-post-adoption-recovery" - register_backend(backend_name, factory, replace=True, factory_api_version=2) - with self.assertRaisesRegex( - BackendFailureError, "recovery_endpoint" - ) as context: - Endpoint("interrupted-v2", backend=backend_name) - recovery = context.exception.recovery_endpoint - self.assertIs(recovery._backend, instances[0]) - self.assertFalse(recovery._backend_receipt.published) - self.assertTrue(recovery._creation_failed) - instances[0].allow_close = True - recovery.close() - self.assertTrue(recovery.closed) - - def test_backend_finalizer_orphans_are_isolated_by_exact_identity(self) -> None: - blocked_ids = {"published-orphan"} - instances = [] - - class OrphanBackend(_ReferenceBackend): - def close(self) -> None: - if self._endpoint_id in blocked_ids: - raise BusyError("injected live backend child") - super().close() - - def factory(*args, adopt_backend, **kwargs) -> None: - del kwargs - backend = OrphanBackend._new_unactivated(*args) - instances.append(backend) - adopt_backend(backend) - backend._activate() - - backend_name = "test-factory-v2-published-orphan" - register_backend(backend_name, factory, replace=True, factory_api_version=2) - abandoned = Endpoint( - "abandoned", backend=backend_name, endpoint_id="published-orphan" - ) - abandoned_ref = weakref.ref(abandoned) - del abandoned - gc.collect() - self.assertIsNone(abandoned_ref()) - self.assertIn( - (backend_name, "published-orphan"), - _api._ORPHANED_BACKEND_RECEIPTS, - ) - - # Auditing the failed orphan does not gate an unrelated endpoint id. - unrelated = Endpoint( - "unrelated", backend=backend_name, endpoint_id="fresh-identity" - ) - unrelated.close() - call_count = len(instances) - with self.assertRaisesRegex( - BackendFailureError, "abandoned published endpoint" - ): - Endpoint("collision", backend=backend_name, endpoint_id="published-orphan") - self.assertEqual(len(instances), call_count) - - blocked_ids.clear() - replacement = Endpoint( - "replacement", backend=backend_name, endpoint_id="published-orphan" - ) - replacement.close() - self.assertNotIn( - (backend_name, "published-orphan"), - _api._ORPHANED_BACKEND_RECEIPTS, - ) - - def test_unpublished_factory_orphan_has_distinct_exact_id_diagnostic(self) -> None: - block_cleanup = True - fail_factory = True - - class OrphanBackend(_ReferenceBackend): - def close(self) -> None: - if block_cleanup: - raise BusyError("injected construction cleanup hold") - super().close() - - def factory(*args, adopt_backend, **kwargs) -> None: - del kwargs - backend = OrphanBackend._new_unactivated(*args) - adopt_backend(backend) - backend._activate() - if fail_factory: - raise RuntimeError("injected post-activation failure") - - backend_name = "test-factory-v2-unpublished-orphan" - register_backend(backend_name, factory, replace=True, factory_api_version=2) - with self.assertRaises(BackendFailureError) as context: - Endpoint( - "unpublished", backend=backend_name, endpoint_id="unpublished-orphan" - ) - recovery = context.exception.recovery_endpoint - recovery_ref = weakref.ref(recovery) - # The test harness retains the raised exception until the test returns; - # sever its documented recovery edge to model an application dropping - # both objects and allow the Endpoint weak finalizer to run now. - context.exception.recovery_endpoint = None - context.exception.__traceback__ = None - context.exception.__cause__ = None - del recovery - del context - gc.collect() - self.assertIsNone(recovery_ref()) - - fail_factory = False - with self.assertRaisesRegex( - BackendFailureError, "unpublished endpoint construction" - ): - Endpoint( - "collision", backend=backend_name, endpoint_id="unpublished-orphan" - ) - block_cleanup = False - replacement = Endpoint( - "replacement", backend=backend_name, endpoint_id="unpublished-orphan" - ) - replacement.close() - - def test_endpoint_close_cut_after_receipt_commit_is_retryable(self) -> None: - endpoint = Endpoint("receipt-close-cut") - original_detach = Endpoint._detach_backend_finalizer - interrupt = True - - def detach_with_cut(owner) -> None: - nonlocal interrupt - if interrupt: - interrupt = False - raise KeyboardInterrupt("after receipt close commit") - original_detach(owner) - - with mock.patch.object(Endpoint, "_detach_backend_finalizer", detach_with_cut): - with self.assertRaisesRegex( - KeyboardInterrupt, "after receipt close commit" - ): - endpoint.close() - self.assertTrue(endpoint._backend_receipt.closed) - self.assertFalse(endpoint.closed) - self.assertTrue(endpoint._backend_finalizer.alive) - endpoint.close() - self.assertTrue(endpoint.closed) - self.assertFalse(endpoint._backend_finalizer.alive) - - def test_endpoint_close_does_not_implicitly_progress(self) -> None: - class NoCloseProgressBackend(_ReferenceBackend): - def progress(self) -> None: - raise AssertionError("close must not inject provider progress") - - backend_name = "test-close-without-progress" - register_backend(backend_name, NoCloseProgressBackend, replace=True) - endpoint = Endpoint("no-close-progress", backend=backend_name) - endpoint.close() - self.assertTrue(endpoint.closed) - - def test_backend_receipt_empty_slots_preserve_unslotted_subclasses(self) -> None: - self.assertEqual(BackendWork.__slots__, ()) - self.assertEqual(BackendRequestSlot.__slots__, ()) - - work = _PendingWork() - slot = _TestRequestSlot() - work.extension_value = "work" - slot.extension_value = "slot" - - self.assertEqual(work.__dict__["extension_value"], "work") - self.assertEqual(slot.__dict__["extension_value"], "slot") - - def test_endpoint_close_retries_without_progress_after_interruption(self) -> None: - endpoint = Endpoint("endpoint", backend="test-interrupt-after-close") - - with self.assertRaises(KeyboardInterrupt): - endpoint.close() - self.assertFalse(endpoint.closed) - self.assertTrue(endpoint._closing) - - endpoint.close() - self.assertTrue(endpoint.closed) - - def test_registration_factory_interruption_rolls_back_adopted_handle(self) -> None: - endpoint = Endpoint("endpoint") - original_register = endpoint._backend.register - - def register_then_interrupt(*args, **kwargs) -> None: - original_register(*args, **kwargs) - raise KeyboardInterrupt - - endpoint._backend.register = register_then_interrupt - with self.assertRaises(KeyboardInterrupt): - endpoint.register(bytearray(4), name="buffer") - - self.assertEqual(endpoint.open_registrations, ()) - self.assertFalse(endpoint._backend._registrations) - endpoint.close() - - def test_failed_registration_rollback_is_publicly_recoverable(self) -> None: - endpoint = Endpoint("endpoint") - original_register = endpoint._backend.register - original_deregister = endpoint._backend.deregister - fail_cleanup = True - - def register_then_interrupt(*args, **kwargs) -> None: - original_register(*args, **kwargs) - raise KeyboardInterrupt - - def deregister_once(handle) -> None: - nonlocal fail_cleanup - if fail_cleanup: - fail_cleanup = False - raise RuntimeError("injected cleanup failure") - original_deregister(handle) - - endpoint._backend.register = register_then_interrupt - endpoint._backend.deregister = deregister_once - with self.assertRaisesRegex(BackendFailureError, "cleanup is pending"): - endpoint.register(bytearray(4), name="buffer") - - (registration,) = endpoint.open_registrations - with self.assertRaisesRegex(BackendFailureError, "cleanup only"): - registration.region() - registration.close() - self.assertEqual(endpoint.open_registrations, ()) - endpoint.close() - - def test_peer_and_plan_factory_interruptions_roll_back(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(4), name="local") - remote = destination.register(bytearray(4), name="remote") - metadata = bytes(destination.export_metadata()) - original_import = source._backend.import_peer - - def import_then_interrupt(*args, **kwargs) -> None: - original_import(*args, **kwargs) - raise KeyboardInterrupt - - source._backend.import_peer = import_then_interrupt - with self.assertRaises(KeyboardInterrupt): - source.import_peer(metadata) - self.assertEqual(source.open_peers, ()) - self.assertFalse(source._backend._peers) - - source._backend.import_peer = original_import - peer = source.import_peer(metadata) - original_prepare = source._backend.prepare - - def prepare_then_interrupt(*args, **kwargs) -> None: - original_prepare(*args, **kwargs) - raise KeyboardInterrupt - - source._backend.prepare = prepare_then_interrupt - with self.assertRaises(KeyboardInterrupt): - source.prepare(local=[local.region()], remote=[peer.region("remote")]) - self.assertEqual(source.open_plans, ()) - self.assertFalse(source._backend._plans) - - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_reference_write_and_metadata_round_trip(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"abcdefgh") - destination_buffer = bytearray(8) - source_registration = source.register(source_buffer, name="blocks") - destination_registration = destination.register( - destination_buffer, name="blocks" - ) - - encoded = bytes(destination.export_metadata()) - self.assertEqual(bytes(EndpointMetadata.from_bytes(encoded)), encoded) - peer = source.import_peer(encoded) - plan = source.prepare( - local=[source_registration.region()], - remote=[peer.region("blocks")], - ) - work = source.write(plan) - self.assertTrue(work.test()) - self.assertEqual(work.state, WorkState.COMPLETED) - self.assertEqual(destination_buffer, source_buffer) - - work.close() - plan.close() - peer.close() - source_registration.close() - destination_registration.close() - source.close() - destination.close() - - def test_reference_read(self) -> None: - sender = Endpoint("sender") - receiver = Endpoint("receiver") - source_buffer = bytearray(b"remote") - destination_buffer = bytearray(6) - source_registration = sender.register(source_buffer, name="weights") - destination_registration = receiver.register(destination_buffer, name="weights") - peer = receiver.import_peer(bytes(sender.export_metadata())) - plan = receiver.prepare( - local=[destination_registration.region()], - remote=[peer.region("weights")], - ) - - work = receiver.read(plan) - work.wait() - self.assertEqual(destination_buffer, source_buffer) - - work.close() - plan.close() - peer.close() - destination_registration.close() - source_registration.close() - receiver.close() - sender.close() - - def test_bound_transfer_full_reuse_and_provider_tokens(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"data") - destination_buffer = bytearray(4) - local = source.register(source_buffer, name="buffer") - remote = destination.register(destination_buffer, name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - tokens: list[object | None] = [] - original_submit = source._backend.submit - - def capture_submit(*args, **kwargs) -> None: - tokens.append(kwargs["reuse_token"]) - original_submit(*args, **kwargs) - - source._backend.submit = capture_submit - first = plan.bind("write") - self.assertIsInstance(first, BoundTransfer) - self.assertEqual(first.op, TransferOp.WRITE) - self.assertEqual(source.open_bound_transfers, (first,)) - - for value in (b"one!", b"two!"): - source_buffer[:] = value - work = first.submit() - work.wait() - self.assertEqual(destination_buffer, value) - work.close() - - second = plan.bind(TransferOp.WRITE) - work = second.submit() - work.close() - reader = plan.bind(TransferOp.READ) - destination_buffer[:] = b"back" - work = reader.submit() - work.wait() - self.assertEqual(source_buffer, b"back") - work.close() - - convenience_start = len(tokens) - for _ in range(2): - work = source.write(plan) - work.close() - - self.assertIs(tokens[0], tokens[1]) - self.assertIsNot(tokens[0], tokens[2]) - self.assertIsNot(tokens[0], tokens[3]) - self.assertIsNot(tokens[2], tokens[3]) - self.assertIs(tokens[convenience_start], tokens[convenience_start + 1]) - self.assertIsNot(tokens[0], tokens[convenience_start]) - - reader.close() - second.close() - first.close() - self.assertEqual(source.open_bound_transfers, ()) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_bound_transfer_canonicalizes_indexed_selection_and_notification( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"a0b0c0d0") - destination_buffer = bytearray(b"--------") - local = source.register(source_buffer, name="source") - remote = destination.register(destination_buffer, name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=2, count=4)], - remote=[peer.region("destination", 0, 1, stride=2, count=4)], - indexed=True, - ) - local_indices = [3, 1] - remote_indices = [0, 2] - notification = bytearray(b"ready") - bound = plan.bind( - TransferOp.WRITE, - local_indices, - remote_indices, - notification, - ) - local_indices[:] = [0] - remote_indices[:] = [1] - notification[:] = b"later" - - work = bound.submit() - work.wait() - self.assertEqual(destination_buffer, b"d---b---") - self.assertEqual( - [item.payload for item in destination.poll_notifications()], - [b"ready"], - ) - - work.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_bound_transfer_enforces_work_plan_endpoint_close_order(self) -> None: - source = Endpoint("source", backend="test-pending") - destination = Endpoint("destination", backend="test-pending") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaisesRegex(BusyError, "bound transfers"): - plan.close() - with self.assertRaises(BusyError): - source.close() - work = bound.submit() - with self.assertRaisesRegex(BusyError, "open work"): - bound.close() - with self.assertRaisesRegex(BusyError, "open work"): - plan.close() - self.assertIn(work, source.open_works) - self.assertIn(bound, source.open_bound_transfers) - - self.assertTrue(work.cancel()) - work.close() - bound.close() - self.assertTrue(bound.closed) - with self.assertRaises(ClosedError): - bound.submit() - self.assertEqual(source.open_bound_transfers, ()) - - plan.close() - with self.assertRaises(ClosedError): - plan.bind(TransferOp.WRITE) - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_bound_transfer_preserves_adoption_interruption_recovery(self) -> None: - source = Endpoint("source", backend="test-adopt-then-interrupt") - destination = Endpoint("destination", backend="test-adopt-then-interrupt") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - work = bound.submit() - self.assertEqual(source.open_works, (work,)) - work.cancel() - self.assertEqual(work.state, WorkState.FAILED) - self.assertIsInstance(work.error, KeyboardInterrupt) - work.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_bound_transfer_rolls_back_interruption_before_adoption(self) -> None: - source = Endpoint("source", backend="test-interrupt") - destination = Endpoint("destination", backend="test-interrupt") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaises(KeyboardInterrupt): - bound.submit() - self.assertEqual(source.open_works, ()) - self.assertFalse(bound._works) - self.assertFalse(plan._submitting) - self.assertFalse(plan._works) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_bound_transfer_close_retry_repairs_post_commit_interruption(self) -> None: - class InterruptingBoundTransfer(BoundTransfer): - __slots__ = ("interrupt_close_commit",) - - def __setattr__(self, name: str, value: object) -> None: - object.__setattr__(self, name, value) - if ( - name == "_closed" - and value is True - and getattr(self, "interrupt_close_commit", False) - ): - object.__setattr__(self, "interrupt_close_commit", False) - raise KeyboardInterrupt - - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - with mock.patch.object(_api, "BoundTransfer", InterruptingBoundTransfer): - bound = plan.bind(TransferOp.WRITE) - bound.interrupt_close_commit = True - - with self.assertRaises(KeyboardInterrupt): - bound.close() - self.assertTrue(bound.closed) - # The parent blockers were removed before the interrupted commit, but - # the retained parent reference lets the same-operation retry finish. - self.assertNotIn(bound, plan._bindings) - self.assertNotIn(bound, source.open_bound_transfers) - self.assertIs(bound._plan, plan) - - bound.close() - self.assertIsNone(bound._plan) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_resource_close_retries_finish_post_commit_epilogues(self) -> None: - class InterruptOnCloseCommit: - __slots__ = () - - def __setattr__(self, name: str, value: object) -> None: - object.__setattr__(self, name, value) - if ( - name == "_closed" - and value is True - and getattr(self, "interrupt_close_commit", False) - ): - object.__setattr__(self, "interrupt_close_commit", False) - raise KeyboardInterrupt - - class InterruptingRegistration(InterruptOnCloseCommit, _api.Registration): - __slots__ = ("interrupt_close_commit",) - - class InterruptingPeer(InterruptOnCloseCommit, _api.Peer): - __slots__ = ("interrupt_close_commit",) - - class InterruptingPlan(InterruptOnCloseCommit, _api.TransferPlan): - __slots__ = ("interrupt_close_commit",) - - class InterruptingWork(InterruptOnCloseCommit, _api.Work): - __slots__ = ("interrupt_close_commit",) - - source = Endpoint("source") - destination = Endpoint("destination") - with mock.patch.object(_api, "Registration", InterruptingRegistration): - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - with mock.patch.object(_api, "Peer", InterruptingPeer): - peer = source.import_peer(bytes(destination.export_metadata())) - with mock.patch.object(_api, "TransferPlan", InterruptingPlan): - plan = source.prepare( - local=[local.region()], remote=[peer.region("buffer")] - ) - with mock.patch.object(_api, "Work", InterruptingWork): - work = source.write(plan) - work.wait() - - for resource in (work, plan, peer, local): - resource.interrupt_close_commit = True - with self.assertRaises(KeyboardInterrupt): - resource.close() - self.assertTrue(resource.closed) - - if resource is work: - self.assertIn(work, source.open_works) - elif resource is plan: - self.assertIn(plan, source.open_plans) - self.assertTrue(plan._origins) - elif resource is peer: - self.assertIn(peer, source.open_peers) - self.assertIsNotNone(peer._handle) - else: - self.assertIn(local, source.open_registrations) - self.assertIsNotNone(local._memory.owner) - - resource.close() - - self.assertEqual(source.open_works, ()) - self.assertEqual(source.open_plans, ()) - self.assertEqual(source.open_peers, ()) - self.assertEqual(source.open_registrations, ()) - remote.close() - source.close() - destination.close() - - def test_bound_host_wait_establishes_ordering_on_every_submit(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local.region() - local_region._device = "cuda:0" - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - source._capabilities = Capabilities( - read=True, - write=True, - manual_progress=True, - background_progress=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - ) - actions: list[str] = [] - - class FakeEvent: - device = torch.device("cuda:0") - - def synchronize(self) -> None: - actions.append("wait") - - with ( - mock.patch.object(torch.cuda, "Event", FakeEvent), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=False - ), - ): - bound = plan.bind( - TransferOp.WRITE, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(), - ) - self.assertEqual(actions, []) - for _ in range(2): - work = bound.submit() - work.close() - self.assertEqual(actions, ["wait", "wait"]) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_backend_request_slot_default_start_and_poll_is_bounded(self) -> None: - class CountingRequestSlot(_TestRequestSlot): - def __init__(self) -> None: - super().__init__() - self.next_state = WorkState.RUNNING - self.poll_count = 0 - - def poll_state(self) -> WorkState: - self.poll_count += 1 - return self.state - - backend_slot = CountingRequestSlot() - self.assertEqual( - backend_slot.start_and_poll(max_polls=3, timeout_ns=None), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.start_count, 1) - self.assertEqual(backend_slot.poll_count, 3) - - backend_slot = CountingRequestSlot() - self.assertEqual( - backend_slot.start_and_poll(max_polls=64, timeout_ns=0), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.poll_count, 1) - backend_slot = CountingRequestSlot() - self.assertEqual( - backend_slot.start_and_poll(max_polls=0, timeout_ns=None), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.start_count, 1) - self.assertEqual(backend_slot.poll_count, 0) - with self.assertRaisesRegex(ValueError, "max_polls"): - backend_slot.start_and_poll(max_polls=-1, timeout_ns=None) - with self.assertRaisesRegex(ValueError, "max_polls"): - backend_slot.start_and_poll(max_polls=1 << 63, timeout_ns=None) - with self.assertRaisesRegex(TypeError, "max_polls"): - backend_slot.start_and_poll(max_polls=True, timeout_ns=None) - with self.assertRaisesRegex(ValueError, "timeout_ns"): - backend_slot.start_and_poll(max_polls=1, timeout_ns=-1) - with self.assertRaisesRegex(ValueError, "timeout_ns"): - backend_slot.start_and_poll(max_polls=1, timeout_ns=1 << 63) - with self.assertRaisesRegex(TypeError, "timeout_ns"): - backend_slot.start_and_poll(max_polls=1, timeout_ns=True) - - backend_slot = CountingRequestSlot() - backend_slot.start = mock.Mock(return_value=object()) - with self.assertRaisesRegex(TypeError, "start must return None"): - backend_slot.start_and_poll(max_polls=1, timeout_ns=None) - - def test_backend_request_slot_default_poll_bounded_contract(self) -> None: - class CountingRequestSlot(_TestRequestSlot): - def __init__(self) -> None: - super().__init__() - self._state = WorkState.RUNNING - self.poll_count = 0 - - def poll_state(self) -> WorkState: - self.poll_count += 1 - assert isinstance(self._state, WorkState) - return self._state - - backend_slot = CountingRequestSlot() - self.assertEqual( - backend_slot.poll_bounded( - max_polls=3, timeout_ns=None, timeout_check_interval=32 - ), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.poll_count, 3) - - backend_slot = CountingRequestSlot() - with mock.patch( - "torch.distributed._transfer._backend.time.monotonic_ns", - side_effect=[10, 10], - ): - self.assertEqual( - backend_slot.poll_bounded( - max_polls=64, timeout_ns=0, timeout_check_interval=32 - ), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.poll_count, 1) - - backend_slot = CountingRequestSlot() - with mock.patch( - "torch.distributed._transfer._backend.time.monotonic_ns", - side_effect=[10, 10, 12], - ): - self.assertEqual( - backend_slot.poll_bounded( - max_polls=64, timeout_ns=1, timeout_check_interval=1 - ), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.poll_count, 2) - - backend_slot = CountingRequestSlot() - with mock.patch( - "torch.distributed._transfer._backend.time.monotonic_ns", - side_effect=[10, 10, 12], - ): - self.assertEqual( - backend_slot.poll_bounded( - max_polls=64, timeout_ns=1, timeout_check_interval=32 - ), - WorkState.RUNNING, - ) - self.assertEqual(backend_slot.poll_count, 32) - - for max_polls, timeout_ns, timeout_check_interval, error_type in ( - (0, None, 32, ValueError), - (True, None, 32, TypeError), - (1, -1, 32, ValueError), - (1 << 63, None, 32, ValueError), - (1, 1 << 63, 32, ValueError), - (1, True, 32, TypeError), - (1, None, 0, ValueError), - (1, None, 1 << 63, ValueError), - (1, None, True, TypeError), - ): - with self.assertRaises(error_type): - backend_slot.poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - - def test_transfer_slot_capability_is_validated_and_opt_in(self) -> None: - with self.assertRaisesRegex(TypeError, "persistent_request_slots"): - Capabilities(persistent_request_slots=1) # type: ignore[arg-type] - with self.assertRaisesRegex(TypeError, "fused_request_slot_start_poll"): - Capabilities( - persistent_request_slots=True, - fused_request_slot_start_poll=1, # type: ignore[arg-type] - ) - with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): - Capabilities(fused_request_slot_start_poll=True) - with self.assertRaisesRegex(TypeError, "fused_request_slot_poll"): - Capabilities( - persistent_request_slots=True, - fused_request_slot_poll=1, # type: ignore[arg-type] - ) - with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): - Capabilities(fused_request_slot_poll=True) - with self.assertRaisesRegex( - TypeError, "deferred_successful_request_slot_recycle" - ): - Capabilities( - persistent_request_slots=True, - deferred_successful_request_slot_recycle=1, # type: ignore[arg-type] - ) - with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): - Capabilities(deferred_successful_request_slot_recycle=True) - with self.assertRaisesRegex(TypeError, "request_slot_notification_overrides"): - Capabilities( - persistent_request_slots=True, - attached_notifications=True, - request_slot_notification_overrides=1, # type: ignore[arg-type] - ) - with self.assertRaisesRegex(TypeError, "fused_prevalidated_submit_poll"): - Capabilities(fused_prevalidated_submit_poll=1) # type: ignore[arg-type] - with self.assertRaisesRegex(TypeError, "fused_work_poll"): - Capabilities(fused_work_poll=1) # type: ignore[arg-type] - with self.assertRaisesRegex(TypeError, "grouped_notification_receive"): - Capabilities(grouped_notification_receive=1) # type: ignore[arg-type] - with self.assertRaisesRegex(ValueError, "requires notifications"): - Capabilities(grouped_notification_receive=True) - with self.assertRaisesRegex(TypeError, "synchronous_notification_send"): - Capabilities(synchronous_notification_send=1) # type: ignore[arg-type] - with self.assertRaisesRegex(ValueError, "requires standalone_notifications"): - Capabilities(synchronous_notification_send=True) - synchronous_notifications = Capabilities( - standalone_notifications=True, - standalone_notification_completion=NotificationCompletion.LOCAL_ACCEPTED, - synchronous_notification_send=True, - ) - self.assertTrue(synchronous_notifications.synchronous_notification_send) - self.assertTrue( - Capabilities( - attached_notifications=True, - grouped_notification_receive=True, - ).grouped_notification_receive - ) - with self.assertRaisesRegex(ValueError, "requires persistent_request_slots"): - Capabilities( - attached_notifications=True, - request_slot_notification_overrides=True, - ) - with self.assertRaisesRegex(ValueError, "requires attached_notifications"): - Capabilities( - persistent_request_slots=True, - request_slot_notification_overrides=True, - ) - poll_only = Capabilities( - persistent_request_slots=True, - fused_request_slot_poll=True, - ) - self.assertTrue(poll_only.fused_request_slot_poll) - self.assertFalse(poll_only.fused_request_slot_start_poll) - self.assertFalse(poll_only.deferred_successful_request_slot_recycle) - - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - self.assertFalse(source.capabilities.persistent_request_slots) - self.assertFalse(source.capabilities.fused_request_slot_start_poll) - self.assertFalse(source.capabilities.fused_request_slot_poll) - self.assertFalse(source.capabilities.deferred_successful_request_slot_recycle) - self.assertFalse(source.capabilities.request_slot_notification_overrides) - self.assertFalse(source.capabilities.fused_prevalidated_submit_poll) - self.assertFalse(source.capabilities.fused_work_poll) - with self.assertRaisesRegex(UnsupportedError, "persistent request slots"): - bound.create_slot() - self.assertEqual(source.open_transfer_slots, ()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_notification_slot_is_explicit_reusable_and_work_free( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-notification") - destination = Endpoint("destination", backend="test-request-slot-notification") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaisesRegex(TypeError, "must be a bool"): - bound.create_slot(request_local_notifications=1) # type: ignore[arg-type] - - fixed_slot = bound.create_slot() - self.assertIs(type(fixed_slot), TransferSlot) - fixed_slot.close() - - slot = bound.create_slot(request_local_notifications=True) - self.assertIs(type(slot), RequestLocalNotificationSlot) - self.assertIsInstance(slot, TransferSlot) - self.assertEqual(source.open_works, ()) - backend_slot = source._backend.request_slots[-1] - - with self.assertRaisesRegex(TypeError, "notification"): - slot.start() # type: ignore[call-arg] - self.assertEqual(slot.generation, 0) - self.assertEqual(backend_slot.start_count, 0) - - released_view = memoryview(bytearray(b"released")) - released_view.release() - for invalid, error_type, message in ( - (b"", ValueError, "must be non-empty"), - (b"x" * (_api._MAX_NOTIFICATION_BYTES + 1), ValueError, "size limit"), - ("wrong", TypeError, "bytes-like"), - (released_view, ValueError, "unavailable"), - ): - with self.assertRaisesRegex(error_type, message): - slot.start(notification=invalid) # type: ignore[arg-type] - self.assertEqual(slot.generation, 0) - self.assertEqual(backend_slot.start_count, 0) - - mutable = bytearray(b"alpha") - slot.start(notification=mutable) - mutable[:] = b"other" - self.assertEqual(backend_slot.notification_starts, [b"alpha"]) - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - - view_owner = bytearray(b"view") - slot.start(notification=memoryview(view_owner)) - view_owner[:] = b"xxxx" - self.assertEqual(backend_slot.notification_starts[-1], b"view") - slot.recycle() - - generation = slot.execute(notification=b"b\x00\xff") - self.assertEqual(generation, 3) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual( - backend_slot.notification_fused_calls, - [(b"b\x00\xff", 64, None)], - ) - - self.assertEqual(slot.execute(notification=None), 4) - self.assertEqual( - backend_slot.notification_fused_calls[-1], - (None, 64, None), - ) - self.assertEqual(source.open_works, ()) - - prior_calls = len(backend_slot.notification_fused_calls) - with self.assertRaisesRegex(ValueError, "must be non-empty"): - slot.execute(notification=b"") - self.assertEqual(slot.generation, 4) - self.assertEqual(len(backend_slot.notification_fused_calls), prior_calls) - - tagged = plan.bind(TransferOp.WRITE, notification=b"fixed") - with self.assertRaisesRegex(ValueError, "notification=None"): - tagged.create_slot(request_local_notifications=True) - tagged.close() - - second = bound.create_slot(request_local_notifications=True) - slot.start(notification=b"one") - second.start(notification=b"two") - self.assertEqual(slot.state, WorkState.COMPLETED) - self.assertEqual(second.state, WorkState.COMPLETED) - slot.recycle() - second.recycle() - second.close() - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_notification_slots_cover_indexed_write_and_read( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-notification") - destination = Endpoint("destination", backend="test-request-slot-notification") - local = source.register(bytearray(b"abcdefgh"), name="buffer") - remote = destination.register(bytearray(8), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 4, stride=4, count=2)], - remote=[peer.region("buffer", 0, 4, stride=4, count=2)], - indexed=True, - ) - - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[1], - remote_indices=[0], - ) - write_bound = selection.bind() - write_slot = write_bound.create_slot(request_local_notifications=True) - self.assertEqual(write_slot.execute(notification=b"indexed"), 1) - - read_bound = plan.bind( - TransferOp.READ, - local_indices=[0], - remote_indices=[1], - ) - read_slot = read_bound.create_slot(request_local_notifications=True) - self.assertEqual(read_slot.execute(notification=b"read"), 1) - - read_slot.close() - read_bound.close() - write_slot.close() - write_bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_notification_slot_is_capability_and_spi_gated( - self, - ) -> None: - for backend, message in ( - ("test-request-slot", "does not support request-local"), - ("test-request-slot-notification-missing", "without adoption"), - ): - source = Endpoint("source", backend=backend) - destination = Endpoint("destination", backend=backend) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region()], remote=[peer.region("buffer")] - ) - bound = plan.bind(TransferOp.WRITE) - with self.assertRaisesRegex( - (UnsupportedError, BackendFailureError), message - ): - bound.create_slot(request_local_notifications=True) - self.assertEqual(source.open_transfer_slots, ()) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_reuses_one_public_resource_and_counts_generations( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - self.assertIsInstance(slot, TransferSlot) - self.assertEqual(source.open_transfer_slots, (slot,)) - self.assertEqual(source.open_works, ()) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot.generation, 0) - - self.assertIsNone(slot.start()) - self.assertTrue(slot.active) - self.assertEqual(slot.generation, 1) - self.assertEqual(slot.state, WorkState.COMPLETED) - self.assertIs(slot.wait(), slot) - backend_slot = source._backend.request_slots[0] - self.assertIs(slot.completion_event, backend_slot.completion) - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.start() - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.close() - - slot.recycle() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot.generation, 1) - self.assertIsNone(slot.start()) - self.assertEqual(slot.generation, 2) - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - self.assertEqual(backend_slot.start_count, 2) - self.assertEqual(backend_slot.recycle_count, 2) - self.assertEqual(source.open_works, ()) - - slot.close() - self.assertTrue(slot.closed) - self.assertEqual(source.open_transfer_slots, ()) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execution_dispatch_is_cached_and_wrapper_owned( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-execution-dispatch") - destination = Endpoint( - "destination", backend="test-request-slot-execution-dispatch" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - dispatch = source._backend.execution_dispatches[0] - self.assertIsNone(slot._fast_execute) - - self.assertEqual( - dispatch.attribute_reads, - {name: 1 for name in _TestExecutionDispatch._OPERATION_NAMES}, - ) - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - self.assertIs(slot.start_and_wait(), slot) - self.assertIs(slot.completion_event, backend_slot.completion) - slot.recycle() - - backend_slot.next_state = WorkState.RUNNING - slot.start() - self.assertEqual(slot.state, WorkState.RUNNING) - self.assertTrue(slot.cancel()) - self.assertEqual(slot.state, WorkState.CANCELLED) - slot.recycle() - - failure = RuntimeError("wrapper-owned diagnostic") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - slot.start() - self.assertIs(slot.error, failure) - slot.recycle() - self.assertEqual( - dispatch.attribute_reads, - {name: 1 for name in _TestExecutionDispatch._OPERATION_NAMES}, - ) - - slot.close() - self.assertEqual(backend_slot.close_count, 1) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execution_dispatch_respects_capability_gating( - self, - ) -> None: - source = Endpoint( - "source", backend="test-request-slot-execution-dispatch-gated" - ) - destination = Endpoint( - "destination", backend="test-request-slot-execution-dispatch-gated" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - dispatch = source._backend.execution_dispatches[0] - - self.assertEqual(dispatch.attribute_reads["start"], 1) - self.assertEqual(dispatch.attribute_reads["poll_state"], 1) - self.assertEqual(dispatch.attribute_reads["cancel"], 1) - self.assertEqual(dispatch.attribute_reads["recycle"], 1) - self.assertEqual(dispatch.attribute_reads["start_and_poll"], 0) - self.assertEqual(dispatch.attribute_reads["poll_bounded"], 0) - slot.start() - slot.recycle() - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_strongly_retains_execution_dispatch_until_close( - self, - ) -> None: - source = Endpoint( - "source", backend="test-request-slot-execution-dispatch-lifetime" - ) - destination = Endpoint( - "destination", backend="test-request-slot-execution-dispatch-lifetime" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - dispatch_ref = source._backend.execution_dispatch_ref - - gc.collect() - self.assertIsNotNone(dispatch_ref()) - self.assertIs(slot._backend_execution_dispatch, dispatch_ref()) - slot.close() - gc.collect() - self.assertIsNone(dispatch_ref()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execution_dispatch_rejects_noncallable_operation( - self, - ) -> None: - for backend in ( - "test-request-slot-execution-dispatch-invalid", - "test-request-slot-execution-dispatch-none-start", - "test-request-slot-execution-dispatch-none-fused", - ): - with self.subTest(backend=backend): - source = Endpoint("source", backend=backend) - destination = Endpoint("destination", backend=backend) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region()], remote=[peer.region("buffer")] - ) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaisesRegex(BackendFailureError, "without adoption"): - bound.create_slot() - self.assertEqual(source.open_transfer_slots, ()) - self.assertEqual(bound._slots, set()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_pending_cancel_failure_and_recycle_contract(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - with self.assertRaisesRegex(BusyError, "no generation"): - slot.recycle() - with self.assertRaisesRegex(BusyError, "no active generation"): - slot.wait() - backend_slot.next_state = WorkState.PENDING - slot.start() - self.assertFalse(slot.test()) - with self.assertRaises(TransferTimeoutError): - slot.wait(0) - with self.assertRaisesRegex(BusyError, "non-terminal"): - slot.recycle() - self.assertTrue(slot.cancel()) - self.assertEqual(slot.state, WorkState.CANCELLED) - self.assertIsInstance(slot.error, TransferCancelledError) - with self.assertRaises(TransferCancelledError): - slot.wait() - slot.recycle() - - failure = RuntimeError("injected slot failure") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - slot.start() - self.assertEqual(slot.state, WorkState.FAILED) - self.assertIs(slot.error, failure) - with self.assertRaises(BackendFailureError) as context: - slot.wait() - self.assertIs(context.exception.__cause__, failure) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_and_wait_fuses_then_falls_back_or_times_out( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - self.assertTrue(source.capabilities.fused_request_slot_start_poll) - self.assertIs(slot.start_and_wait(), slot) - self.assertEqual(backend_slot.fused_calls, [(64, None)]) - self.assertEqual(backend_slot.poll_count, 0) - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - - self.assertIs(slot.start_and_wait(1 << 80), slot) - self.assertEqual(backend_slot.fused_calls[-1], (64, (1 << 63) - 1)) - slot.recycle() - - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.COMPLETED] - self.assertIs(slot.start_and_wait(1), slot) - max_polls, timeout_ns = backend_slot.fused_calls[-1] - self.assertEqual(max_polls, 64) - self.assertIsInstance(timeout_ns, int) - self.assertGreaterEqual(timeout_ns, 0) - self.assertLessEqual(timeout_ns, 1_000_000_000) - self.assertEqual(backend_slot.poll_count, 1) - slot.recycle() - - backend_slot.next_state = WorkState.RUNNING - with self.assertRaises(TransferTimeoutError): - slot.start_and_wait(0) - self.assertEqual(backend_slot.fused_calls[-1], (64, 0)) - self.assertTrue(slot.active) - self.assertEqual(backend_slot.poll_count, 1) - self.assertEqual(slot.state, WorkState.RUNNING) - self.assertTrue(slot.cancel()) - slot.recycle() - - generation = slot.generation - for timeout, error_type in ( - (True, TypeError), - ("one", TypeError), - (-1, ValueError), - (float("inf"), ValueError), - ): - with self.assertRaises(error_type): - slot.start_and_wait(timeout) # type: ignore[arg-type] - self.assertEqual(slot.generation, generation) - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_and_wait_uses_bounded_poll_chunks(self) -> None: - source = Endpoint( - "source", backend="test-request-slot-fused", thread_mode="single" - ) - destination = Endpoint( - "destination", backend="test-request-slot-fused", thread_mode="single" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [ - WorkState.RUNNING, - WorkState.RUNNING, - WorkState.COMPLETED, - ] - - with mock.patch.object( - source, - "_backend_request_slot_poll_bounded", - side_effect=AssertionError( - "SINGLE bounded polling must use direct cached dispatch" - ), - ): - self.assertIs(slot.start_and_wait(), slot) - self.assertEqual( - backend_slot.bounded_calls, - [ - ( - _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, - _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, - ) - ] - * 3, - ) - self.assertEqual(backend_slot.poll_count, 0) - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_poll_only_continuation_does_not_require_start_fusion( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-poll-only") - destination = Endpoint("destination", backend="test-request-slot-poll-only") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - - self.assertFalse(source.capabilities.fused_request_slot_start_poll) - self.assertTrue(source.capabilities.fused_request_slot_poll) - self.assertIs(slot.start_and_wait(), slot) - self.assertEqual(backend_slot.start_count, 1) - self.assertEqual(backend_slot.fused_calls, []) - self.assertEqual( - backend_slot.bounded_calls, - [ - ( - _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, - _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, - ) - ], - ) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_bounded_poll_timeout_leaves_generation_active( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.RUNNING] - - with mock.patch( - "torch.distributed._transfer._api.time.monotonic_ns", - side_effect=[0, 0, 0, 999_950_000, 1_000_000_000], - ): - with self.assertRaises(TransferTimeoutError): - slot.start_and_wait(1) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.RUNNING) - self.assertEqual( - backend_slot.bounded_calls, - [ - ( - _api._FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - 50_000, - _api._FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL, - ) - ], - ) - self.assertTrue(slot.cancel()) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_bounded_poll_validates_provider_results(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - for result, message in ( - (None, "invalid transfer slot state"), - ("running", "invalid transfer slot state"), - (WorkState.PENDING, "regressed"), - ): - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [result] - with self.assertRaisesRegex(BackendFailureError, message): - slot.start_and_wait() - self.assertTrue(slot.active) - self.assertEqual(slot._state, WorkState.RUNNING) - backend_slot._state = WorkState.RUNNING - self.assertTrue(slot.cancel()) - slot.recycle() - - provider_error = RuntimeError("injected bounded-poll failure") - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_error = provider_error - with self.assertRaises(BackendFailureError) as context: - slot.start_and_wait() - self.assertIs(context.exception.__cause__, provider_error) - self.assertTrue(slot.active) - self.assertEqual(slot._state, WorkState.RUNNING) - backend_slot.bounded_error = None - backend_slot._state = WorkState.RUNNING - self.assertTrue(slot.cancel()) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_cancel_can_acquire_between_bounded_poll_chunks( - self, - ) -> None: - source = Endpoint( - "source", backend="test-request-slot-fused", thread_mode="multiple" - ) - destination = Endpoint( - "destination", backend="test-request-slot-fused", thread_mode="multiple" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.RUNNING] - between_chunks = threading.Event() - resume_waiter = threading.Event() - waiter_errors: list[BaseException] = [] - original_progress = Endpoint.progress - - def pause_between_chunks(self): - if self is source and not between_chunks.is_set(): - between_chunks.set() - if not resume_waiter.wait(2): - raise RuntimeError("bounded-poll cancellation test timed out") - return original_progress(self) - - def start_and_wait() -> None: - try: - slot.start_and_wait() - except BaseException as error: - waiter_errors.append(error) - - with mock.patch.object(Endpoint, "progress", pause_between_chunks): - waiter = threading.Thread(target=start_and_wait) - waiter.start() - self.assertTrue(between_chunks.wait(1)) - self.assertTrue(slot.cancel()) - resume_waiter.set() - waiter.join(2) - - self.assertFalse(waiter.is_alive()) - self.assertEqual(len(waiter_errors), 1) - self.assertIsInstance(waiter_errors[0], TransferCancelledError) - self.assertEqual(len(backend_slot.bounded_calls), 1) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_bounded_wait_is_pinned_to_started_generation(self) -> None: - source = Endpoint( - "source", backend="test-request-slot-fused", thread_mode="multiple" - ) - destination = Endpoint( - "destination", backend="test-request-slot-fused", thread_mode="multiple" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - entered = threading.Event() - release = threading.Event() - errors: list[BaseException] = [] - original_wait_bounded = TransferSlot._wait_bounded_until - - def pause_before_bounded_wait(self, *args, **kwargs): - entered.set() - if not release.wait(2): - raise RuntimeError("bounded-poll generation test timed out") - return original_wait_bounded(self, *args, **kwargs) - - def run_waiter() -> None: - try: - slot.start_and_wait() - except BaseException as error: - errors.append(error) - - with mock.patch.object( - TransferSlot, "_wait_bounded_until", pause_before_bounded_wait - ): - waiter = threading.Thread(target=run_waiter) - waiter.start() - self.assertTrue(entered.wait(1)) - backend_slot._state = WorkState.COMPLETED - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - backend_slot.next_state = WorkState.COMPLETED - slot.start() - replacement_generation = slot.generation - release.set() - waiter.join(2) - - self.assertFalse(waiter.is_alive()) - self.assertEqual(len(errors), 1) - self.assertIsInstance(errors[0], BusyError) - self.assertIn("changed while waiting", str(errors[0])) - self.assertEqual(slot.generation, replacement_generation) - self.assertEqual(backend_slot.bounded_calls, []) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_and_wait_validates_fused_provider_results( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - failure = RuntimeError("fused provider failure") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - with self.assertRaises(BackendFailureError) as context: - slot.start_and_wait() - self.assertIs(context.exception.__cause__, failure) - self.assertIs(slot.error, failure) - slot.recycle() - - backend_slot.next_state = WorkState.CANCELLED - backend_slot.next_error = None - with self.assertRaises(TransferCancelledError): - slot.start_and_wait() - self.assertIsInstance(slot.error, TransferCancelledError) - slot.recycle() - - for result, message in ( - (None, "without acknowledging"), - ("completed", "invalid fused request-slot state"), - ): - backend_slot.next_state = WorkState.COMPLETED - backend_slot.fused_result = result - with self.assertRaisesRegex(BackendFailureError, message): - slot.start_and_wait() - self.assertEqual(slot.state, WorkState.FAILED) - slot.recycle() - - backend_slot.fused_result = _RETURN_CURRENT_SLOT_STATE - backend_slot.fused_error = KeyboardInterrupt() - backend_slot.next_state = WorkState.COMPLETED - with self.assertRaises(BackendFailureError) as context: - slot.start_and_wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertEqual(slot.state, WorkState.FAILED) - slot.recycle() - - backend_slot.fused_error = None - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.PENDING] - with self.assertRaisesRegex(BackendFailureError, "regressed"): - slot.start_and_wait(1) - self.assertTrue(slot.active) - self.assertEqual(slot._state, WorkState.RUNNING) - backend_slot._state = WorkState.RUNNING - self.assertTrue(slot.cancel()) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_and_wait_fallback_ignores_fused_spi(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.start_and_poll = mock.Mock( - side_effect=AssertionError("fused SPI must not be entered") - ) - backend_slot.poll_bounded = mock.Mock( - side_effect=AssertionError("bounded-poll SPI must not be entered") - ) - - self.assertFalse(source.capabilities.fused_request_slot_start_poll) - self.assertFalse(source.capabilities.fused_request_slot_poll) - self.assertIs(slot.start_and_wait(), slot) - backend_slot.start_and_poll.assert_not_called() - backend_slot.poll_bounded.assert_not_called() - self.assertEqual(backend_slot.start_count, 1) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_source_compatible_fallback_recycles(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - self.assertFalse(source.capabilities.deferred_successful_request_slot_recycle) - self.assertEqual(slot.execute(), 1) - self.assertEqual(slot.generation, 1) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(backend_slot.start_count, 1) - self.assertEqual(backend_slot.recycle_count, 1) - with self.assertRaisesRegex(BusyError, "no active generation"): - _ = slot.completion_event - with self.assertRaisesRegex(BusyError, "no generation"): - slot.recycle() - - generation = slot.generation - for timeout, error_type in ( - (True, TypeError), - ("one", TypeError), - (-1, ValueError), - (float("inf"), ValueError), - ): - with self.assertRaises(error_type): - slot.execute(timeout) # type: ignore[arg-type] - self.assertEqual(slot.generation, generation) - self.assertFalse(slot.active) - - self.assertEqual(slot.execute(), 2) - self.assertEqual(backend_slot.recycle_count, 2) - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_deferred_recycle_preserves_explicit_api( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-deferred-recycle") - destination = Endpoint( - "destination", backend="test-request-slot-deferred-recycle" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - self.assertTrue(source.capabilities.deferred_successful_request_slot_recycle) - self.assertEqual(slot.execute(), 1) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertFalse(slot.test()) - self.assertIsNone(slot.error) - self.assertEqual(backend_slot.recycle_count, 0) - # The test provider models the retained, observable completion receipt. - self.assertEqual(backend_slot.state, WorkState.COMPLETED) - with self.assertRaisesRegex(BusyError, "no active generation"): - _ = slot.completion_event - - self.assertIs(slot.start_and_wait(), slot) - self.assertEqual(slot.generation, 2) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.COMPLETED) - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.start() - with self.assertRaisesRegex(BusyError, "recycle it before close"): - slot.close() - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - # Close after execute is legal even though the provider still owns its - # deferred successful receipt. - self.assertEqual(slot.execute(), 3) - self.assertEqual(backend_slot.recycle_count, 1) - slot.close() - self.assertTrue(slot.closed) - self.assertEqual(backend_slot.close_count, 1) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_deferred_fused_immediate_success(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - with ( - mock.patch.object( - source, - "_backend_request_slot_start_and_poll", - side_effect=AssertionError("SINGLE execute must use cached dispatch"), - ), - mock.patch.object( - _api, - "_transfer_slot_deadline_ns", - side_effect=AssertionError("timeout=None must bypass validation"), - ), - mock.patch.object( - TransferSlot, - "_check_access", - side_effect=AssertionError("fast execute must inline access checks"), - ), - mock.patch.object( - TransferSlot, - "_begin_generation_unlocked", - side_effect=AssertionError("fast execute must inline generation begin"), - ), - mock.patch.object( - TransferSlot, - "_start_and_poll_unlocked", - side_effect=AssertionError("fast execute must inline fused start"), - ), - mock.patch.object( - TransferSlot, - "_observe_backend_state_unlocked", - side_effect=AssertionError("fast execute must inline completion"), - ), - ): - self.assertEqual(slot.execute(), 1) - self.assertEqual(slot.execute(), 2) - self.assertEqual(backend_slot.fused_calls, [(64, None), (64, None)]) - self.assertEqual(backend_slot.poll_count, 0) - self.assertEqual(backend_slot.recycle_count, 0) - self.assertEqual(slot.generation, 2) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_exact_running_uses_direct_continuation( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - observe = TransferSlot._observe_backend_state_unlocked - - def reject_initial_running(self, state): - if state is WorkState.RUNNING: - raise AssertionError("initial RUNNING must use direct continuation") - return observe(self, state) - - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - with mock.patch.object( - TransferSlot, - "_observe_backend_state_unlocked", - reject_initial_running, - ): - self.assertEqual(slot.execute(), 1) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(len(backend_slot.bounded_calls), 1) - - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.PENDING] - with ( - mock.patch.object( - TransferSlot, - "_observe_backend_state_unlocked", - reject_initial_running, - ), - self.assertRaisesRegex(BackendFailureError, "regressed"), - ): - slot.execute() - self.assertTrue(slot.active) - self.assertTrue(slot._backend_generation_active) - self.assertIs(slot._state, WorkState.RUNNING) - backend_slot._state = WorkState.RUNNING - self.assertTrue(slot.cancel()) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_test_recycles_only_success(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - with mock.patch.object( - TransferSlot, - "_start_and_poll_unlocked", - side_effect=AssertionError("exact start/test must use cached dispatch"), - ): - self.assertTrue(slot.start_and_test()) - self.assertFalse(slot.active) - self.assertEqual(slot.generation, 1) - self.assertEqual(backend_slot.fused_calls, [(0, None)]) - self.assertEqual(backend_slot.recycle_count, 0) - - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.RUNNING, WorkState.COMPLETED] - with mock.patch.object( - TransferSlot, - "_observe_backend_state_unlocked", - side_effect=AssertionError( - "exact RUNNING start/test must not reclassify provider state" - ), - ): - self.assertFalse(slot.start_and_test()) - self.assertTrue(slot.active) - with ( - mock.patch.object( - TransferSlot, - "_check_access", - side_effect=AssertionError("exact continuation must inline access"), - ), - mock.patch.object( - TransferSlot, - "_refresh_unlocked", - side_effect=AssertionError("exact continuation must poll directly"), - ), - mock.patch.object( - TransferSlot, - "_finish_execute_success_unlocked", - side_effect=AssertionError("exact continuation must finish inline"), - ), - mock.patch.object( - TransferSlot, - "_observe_backend_state_unlocked", - side_effect=AssertionError( - "exact continuation must inline common provider states" - ), - ), - ): - self.assertFalse(slot.test_and_recycle()) - self.assertTrue(slot.test_and_recycle()) - self.assertFalse(slot.active) - self.assertEqual(backend_slot.recycle_count, 0) - with self.assertRaisesRegex(BusyError, "no active generation"): - slot.test_and_recycle() - - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = RuntimeError("injected failure") - with self.assertRaises(BackendFailureError): - slot.start_and_test() - self.assertTrue(slot.active) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_cached_success_recycles_without_another_poll(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.COMPLETED] - self.assertFalse(slot.start_and_test()) - self.assertTrue(slot.test()) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.COMPLETED) - polls_before_recycle = backend_slot.poll_count - with mock.patch.object( - slot, - "_backend_poll_state", - side_effect=AssertionError("cached completion must not poll again"), - ): - self.assertTrue(slot.test_and_recycle()) - self.assertEqual(backend_slot.poll_count, polls_before_recycle) - self.assertEqual(backend_slot.recycle_count, 0) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_serialized_fallback_start_test_is_one_transaction( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-blocking-observation", - thread_mode="serialized", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-blocking-observation", - thread_mode="serialized", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.block_next_observation = True - results: dict[str, bool] = {} - errors: list[BaseException] = [] - second_entered = threading.Event() - second_done = threading.Event() - - def start_and_test(label: str) -> None: - if label == "second": - second_entered.set() - try: - results[label] = slot.start_and_test() - except BaseException as error: - errors.append(error) - finally: - if label == "second": - second_done.set() - - first = threading.Thread(target=start_and_test, args=("first",)) - second = threading.Thread(target=start_and_test, args=("second",)) - first.start() - self.assertTrue(backend_slot.observation_entered.wait(1)) - second.start() - self.assertTrue(second_entered.wait(1)) - try: - self.assertFalse(second_done.wait(0.05)) - self.assertEqual(backend_slot.start_count, 1) - self.assertEqual(slot.generation, 1) - self.assertTrue(slot.active) - finally: - backend_slot.observation_release.set() - first.join(2) - second.join(2) - - self.assertFalse(first.is_alive()) - self.assertFalse(second.is_alive()) - self.assertEqual(errors, []) - self.assertEqual(results, {"first": True, "second": True}) - self.assertEqual(backend_slot.start_count, 2) - self.assertEqual(backend_slot.recycle_count, 2) - self.assertEqual(slot.generation, 2) - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_test_retains_cancel_and_failure(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - backend_slot.next_state = WorkState.RUNNING - self.assertFalse(slot.start_and_test()) - self.assertTrue(slot.cancel()) - with self.assertRaises(TransferCancelledError): - slot.test_and_recycle() - self.assertTrue(slot.active) - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.start_and_test() - with self.assertRaisesRegex(BusyError, "recycle it before close"): - slot.close() - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - failure = RuntimeError("retained start/test failure") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - with self.assertRaises(BackendFailureError) as context: - slot.start_and_test() - self.assertIs(context.exception.__cause__, failure) - self.assertTrue(slot.active) - with self.assertRaises(BackendFailureError) as context: - slot.test_and_recycle() - self.assertIs(context.exception.__cause__, failure) - self.assertTrue(slot.active) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 2) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_test_deferred_success_closes_without_recycle( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - self.assertTrue(slot.start_and_test()) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(backend_slot.recycle_count, 0) - slot.close() - self.assertTrue(slot.closed) - self.assertEqual(backend_slot.recycle_count, 0) - self.assertEqual(backend_slot.close_count, 1) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_test_interruption_boundaries_are_recoverable( - self, - ) -> None: - class InterruptingIdleCommitTransferSlot(TransferSlot): - __slots__ = ("interrupt_idle_commit",) - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "TransferSlot", InterruptingIdleCommitTransferSlot - ): - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - backend_slot.fused_error = KeyboardInterrupt() - with self.assertRaises(BackendFailureError) as context: - slot.start_and_test() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.FAILED) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - backend_slot.fused_error = None - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.start_and_test() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot.generation, 2) - self.assertEqual(backend_slot.recycle_count, 1) - self.assertTrue(slot.start_and_test()) - self.assertEqual(slot.generation, 3) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_slot_start_test_snapshots_notification(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-notification", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-notification", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot(request_local_notifications=True) - backend_slot = source._backend.request_slots[0] - - payload = bytearray(b"A\x00\xff") - with mock.patch.object( - RequestLocalNotificationSlot, - "_start_and_poll_with_notification_unlocked", - side_effect=AssertionError( - "exact notification start/test must use cached dispatch" - ), - ): - self.assertTrue(slot.start_and_test(notification=payload)) - payload[:] = b"BAD" - self.assertEqual( - backend_slot.notification_fused_calls, - [(b"A\x00\xff", 0, None)], - ) - self.assertTrue(slot.start_and_test(notification=None)) - self.assertEqual( - backend_slot.notification_fused_calls[-1], - (None, 0, None), - ) - self.assertFalse(slot.active) - self.assertEqual(backend_slot.recycle_count, 0) - - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.COMPLETED] - with mock.patch.object( - RequestLocalNotificationSlot, - "_observe_backend_state_unlocked", - side_effect=AssertionError( - "exact notification RUNNING start/test must not reclassify state" - ), - ): - self.assertFalse(slot.start_and_test(notification=b"pending")) - self.assertTrue(slot.active) - self.assertTrue(slot.test_and_recycle()) - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_slot_execute_exact_running_uses_direct_continuation( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-notification", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-notification", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot(request_local_notifications=True) - backend_slot = source._backend.request_slots[0] - observe = RequestLocalNotificationSlot._observe_backend_state_unlocked - - def reject_initial_running(self, state): - if state is WorkState.RUNNING: - raise AssertionError("initial RUNNING must use direct continuation") - return observe(self, state) - - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - with mock.patch.object( - RequestLocalNotificationSlot, - "_observe_backend_state_unlocked", - reject_initial_running, - ): - self.assertEqual(slot.execute(notification=b"ready"), 1) - self.assertEqual( - backend_slot.notification_fused_calls, - [(b"ready", _api._FUSED_REQUEST_SLOT_MAX_POLLS, None)], - ) - self.assertEqual(len(backend_slot.bounded_calls), 1) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertIs( - slot._notification_fast_execute, - slot._backend_start_and_poll_with_notification, - ) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_slot_interrupted_completion_cannot_reuse_stale_state( - self, - ) -> None: - class InterruptingNotificationSlot(RequestLocalNotificationSlot): - __slots__ = ("interrupt_idle_commit",) - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint( - "source", backend="test-request-slot-notification", thread_mode="single" - ) - destination = Endpoint( - "destination", - backend="test-request-slot-notification", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "RequestLocalNotificationSlot", InterruptingNotificationSlot - ): - slot = bound.create_slot(request_local_notifications=True) - backend_slot = source._backend.request_slots[0] - notification_fast_execute = slot._notification_fast_execute - self.assertIsNotNone(notification_fast_execute) - - backend_slot.next_state = WorkState.RUNNING - self.assertFalse(slot.start_and_test(notification=b"first")) - backend_slot.poll_states = [WorkState.COMPLETED] - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.test_and_recycle() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertIsNone(slot._state) - self.assertIsNone(slot._start_error) - self.assertIs( - slot._notification_fast_execute, - notification_fast_execute, - ) - - prior_poll_count = backend_slot.poll_count - backend_slot.next_state = WorkState.RUNNING - backend_slot.fused_error = KeyboardInterrupt() - backend_slot.poll_states = [WorkState.COMPLETED] - with self.assertRaises(BackendFailureError) as context: - slot.start_and_test(notification=b"second") - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertEqual(backend_slot.poll_count, prior_poll_count + 1) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.FAILED) - - slot.recycle() - self.assertFalse(slot.active) - self.assertIs( - slot._notification_fast_execute, - slot._backend_start_and_poll_with_notification, - ) - backend_slot.fused_error = None - backend_slot.next_state = WorkState.COMPLETED - self.assertTrue(slot.start_and_test(notification=None)) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_request_local_slot_interrupted_bounded_execute_uses_safe_fallback( - self, - ) -> None: - class InterruptingNotificationSlot(RequestLocalNotificationSlot): - __slots__ = ("interrupt_idle_commit",) - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint( - "source", backend="test-request-slot-notification", thread_mode="single" - ) - destination = Endpoint( - "destination", - backend="test-request-slot-notification", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "RequestLocalNotificationSlot", InterruptingNotificationSlot - ): - slot = bound.create_slot(request_local_notifications=True) - backend_slot = source._backend.request_slots[0] - - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.execute(notification=b"bounded") - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot._state, WorkState.COMPLETED) - self.assertIsNone(slot._notification_fast_execute) - - backend_slot.next_state = WorkState.COMPLETED - self.assertEqual(slot.execute(notification=b"fallback"), 2) - self.assertFalse(slot.active) - self.assertIsNone(slot._state) - self.assertIs( - slot._notification_fast_execute, - slot._backend_start_and_poll_with_notification, - ) - - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = RuntimeError("notification failure") - with self.assertRaises(BackendFailureError): - slot.start_and_test(notification=b"failure") - self.assertIsNotNone(slot._notification_fast_execute) - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.recycle() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot._state, WorkState.FAILED) - self.assertIsNone(slot._notification_fast_execute) - - backend_slot.next_state = WorkState.COMPLETED - backend_slot.next_error = None - self.assertTrue(slot.start_and_test(notification=None)) - self.assertIs( - slot._notification_fast_execute, - slot._backend_start_and_poll_with_notification, - ) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_single_fused_fast_path_preserves_guards( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - wrong_thread_errors: list[BaseException] = [] - - def execute_from_wrong_thread() -> None: - try: - slot.execute() - except BaseException as error: - wrong_thread_errors.append(error) - - wrong_thread = threading.Thread(target=execute_from_wrong_thread) - wrong_thread.start() - wrong_thread.join(1) - self.assertFalse(wrong_thread.is_alive()) - self.assertEqual(len(wrong_thread_errors), 1) - self.assertIsInstance(wrong_thread_errors[0], UnsupportedError) - self.assertEqual(slot.generation, 0) - self.assertFalse(slot.active) - self.assertEqual(backend_slot.fused_calls, []) - - backend_slot.next_state = WorkState.COMPLETED - backend_slot.fused_error = KeyboardInterrupt() - with self.assertRaises(BackendFailureError) as context: - slot.execute() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertEqual(slot.generation, 1) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.FAILED) - slot.recycle() - - backend_slot.fused_error = None - slot.start() - fused_call_count = len(backend_slot.fused_calls) - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.execute() - self.assertEqual(len(backend_slot.fused_calls), fused_call_count) - self.assertEqual(slot.generation, 2) - slot.recycle() - - self.assertEqual(slot.execute(), 3) - slot.close() - with self.assertRaises(ClosedError): - slot.execute() - self.assertEqual(slot.generation, 3) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_bounded_success_is_atomic_in_multiple( - self, - ) -> None: - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="multiple", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="multiple", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - finish_entered = threading.Event() - finish_release = threading.Event() - observer_done = threading.Event() - execute_results: list[int] = [] - errors: list[BaseException] = [] - observed_states: list[WorkState | None] = [] - original_finish = TransferSlot._finish_execute_success_unlocked - - def pause_in_finish(self, expected_generation): - finish_entered.set() - if not finish_release.wait(2): - raise RuntimeError("execute atomicity test timed out") - return original_finish(self, expected_generation) - - def run_execute() -> None: - try: - execute_results.append(slot.execute()) - except BaseException as error: - errors.append(error) - - def observe_state() -> None: - try: - observed_states.append(slot.state) - except BaseException as error: - errors.append(error) - finally: - observer_done.set() - - with mock.patch.object( - TransferSlot, "_finish_execute_success_unlocked", pause_in_finish - ): - executor = threading.Thread(target=run_execute) - executor.start() - self.assertTrue(finish_entered.wait(1)) - observer = threading.Thread(target=observe_state) - observer.start() - self.assertFalse(observer_done.wait(0.05)) - finish_release.set() - executor.join(2) - observer.join(2) - - self.assertFalse(executor.is_alive()) - self.assertFalse(observer.is_alive()) - self.assertEqual(errors, []) - self.assertEqual(execute_results, [1]) - self.assertEqual(observed_states, [None]) - self.assertEqual(len(backend_slot.bounded_calls), 1) - self.assertEqual(backend_slot.recycle_count, 0) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_retains_unsuccessful_generations(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused-deferred-recycle") - destination = Endpoint( - "destination", backend="test-request-slot-fused-deferred-recycle" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - failure = RuntimeError("execute provider failure") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - with self.assertRaises(BackendFailureError) as context: - slot.execute() - self.assertIs(context.exception.__cause__, failure) - self.assertTrue(slot.active) - self.assertEqual(slot.generation, 1) - slot.recycle() - - backend_slot.next_state = WorkState.CANCELLED - backend_slot.next_error = None - with self.assertRaises(TransferCancelledError): - slot.execute() - self.assertTrue(slot.active) - self.assertEqual(slot.generation, 2) - slot.recycle() - - backend_slot.next_state = WorkState.RUNNING - with self.assertRaises(TransferTimeoutError): - slot.execute(0) - self.assertTrue(slot.active) - self.assertEqual(slot.generation, 3) - self.assertEqual(slot.state, WorkState.RUNNING) - self.assertTrue(slot.cancel()) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 3) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_execute_interruption_preserves_a_safe_boundary( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-fused-deferred-recycle") - destination = Endpoint( - "destination", backend="test-request-slot-fused-deferred-recycle" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - original_finish = TransferSlot._finish_execute_success_unlocked - - def interrupt_before_finish(self, expected_generation): - raise KeyboardInterrupt - - with mock.patch.object( - TransferSlot, - "_finish_execute_success_unlocked", - interrupt_before_finish, - ): - with self.assertRaises(KeyboardInterrupt): - slot.execute() - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - def interrupt_after_finish(self, expected_generation): - original_finish(self, expected_generation) - raise KeyboardInterrupt - - with mock.patch.object( - TransferSlot, - "_finish_execute_success_unlocked", - interrupt_after_finish, - ): - with self.assertRaises(KeyboardInterrupt): - slot.execute() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot.generation, 2) - self.assertEqual(backend_slot.recycle_count, 1) - - self.assertEqual(slot.execute(), 3) - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_deferred_execute_idle_boundaries_are_recoverable( - self, - ) -> None: - class InterruptingIdleCommitTransferSlot(TransferSlot): - __slots__ = ("interrupt_fast_restore", "interrupt_idle_commit") - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_fast_restore = False - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - if ( - name == "_fast_execute" - and value is not None - and getattr(self, "interrupt_fast_restore", False) - ): - object.__setattr__(self, "interrupt_fast_restore", False) - raise KeyboardInterrupt - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "TransferSlot", InterruptingIdleCommitTransferSlot - ): - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - fast_execute = slot._fast_execute - self.assertIs(fast_execute, slot._backend_start_and_poll) - self.assertTrue(slot._caller_ready) - - def interrupt_before_provider_entry(**_kwargs) -> WorkState: - raise KeyboardInterrupt - - slot._fast_execute = interrupt_before_provider_entry - with self.assertRaises(BackendFailureError) as context: - slot.execute() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertTrue(slot.active) - self.assertEqual(slot.state, WorkState.FAILED) - self.assertEqual(backend_slot.fused_calls, []) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 0) - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.execute() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertIsNone(slot.error) - self.assertIsNone(slot._state) - self.assertEqual(slot.generation, 2) - self.assertEqual(backend_slot.state, WorkState.COMPLETED) - self.assertIs(slot._fast_execute, fast_execute) - self.assertEqual(slot.execute(), 3) - self.assertEqual(slot.execute(), 4) - self.assertIsNone(slot._state) - self.assertEqual(backend_slot.recycle_count, 0) - - # A timed call uses the generic success finalizer. It disables exact - # execution before public idle, so an interrupt at that commit leaves a - # dirty but fast-ineligible slot that the next generic call repairs. - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.execute(1.0) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(slot._state, WorkState.COMPLETED) - self.assertIsNone(slot._fast_execute) - - # Publishing exact eligibility is the last private-idle store. An - # interrupt just before it leaves a completely clean generic fallback; - # one successful fallback restores the exact path for its successor. - slot.interrupt_fast_restore = True - with self.assertRaises(KeyboardInterrupt): - slot.execute(1.0) - self.assertFalse(slot.active) - self.assertIsNone(slot._state) - self.assertIsNone(slot._error) - self.assertIsNone(slot._start_error) - self.assertFalse(slot._backend_generation_active) - self.assertIsNone(slot._fast_execute) - self.assertEqual(slot.execute(), 7) - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - self.assertEqual(slot.execute(), 8) - - slot.close() - self.assertIsNone(slot._fast_execute) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_dirty_idle_repairs_ambiguous_failure_before_begin( - self, - ) -> None: - class InterruptingIdleCommitTransferSlot(TransferSlot): - __slots__ = ("interrupt_fast_restore", "interrupt_idle_commit") - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_fast_restore = False - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - if ( - name == "_fast_execute" - and value is not None - and getattr(self, "interrupt_fast_restore", False) - ): - object.__setattr__(self, "interrupt_fast_restore", False) - raise KeyboardInterrupt - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint( - "source", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-fused-deferred-recycle", - thread_mode="single", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "TransferSlot", InterruptingIdleCommitTransferSlot - ): - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - backend_slot.next_state = WorkState.COMPLETED - backend_slot.fused_error = KeyboardInterrupt() - with self.assertRaises(BackendFailureError): - slot.execute() - self.assertEqual(slot.state, WorkState.FAILED) - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.recycle() - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertIsNone(slot.error) - self.assertEqual(slot._state, WorkState.FAILED) - self.assertIsNotNone(slot._start_error) - self.assertIsNone(slot._fast_execute) - - backend_slot.fused_error = None - backend_slot.next_state = WorkState.RUNNING - backend_slot.bounded_results = [WorkState.COMPLETED] - self.assertEqual(slot.execute(), 2) - self.assertFalse(slot.active) - self.assertIsNone(slot.error) - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - - # CANCELLED is never folded. Its explicit recycle follows the same - # disable-clean-restore protocol; interruption before the last restore - # leaves a clean generic fallback and cannot expose stale diagnostics. - backend_slot.next_state = WorkState.CANCELLED - with self.assertRaises(TransferCancelledError): - slot.execute() - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - slot.interrupt_fast_restore = True - with self.assertRaises(KeyboardInterrupt): - slot.recycle() - self.assertFalse(slot.active) - self.assertIsNone(slot._state) - self.assertIsNone(slot._error) - self.assertIsNone(slot._start_error) - self.assertFalse(slot._backend_generation_active) - self.assertIsNone(slot._fast_execute) - - backend_slot.next_state = WorkState.COMPLETED - self.assertEqual(slot.execute(), 4) - self.assertIs(slot._fast_execute, slot._backend_start_and_poll) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_fused_fallback_releases_slot_lock(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-start-fused-only", - thread_mode="multiple", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-start-fused-only", - thread_mode="multiple", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.start_entered = threading.Event() - errors: list[BaseException] = [] - - def start_and_wait() -> None: - try: - slot.start_and_wait(2) - except BaseException as error: - errors.append(error) - - waiter = threading.Thread(target=start_and_wait) - waiter.start() - self.assertTrue(backend_slot.start_entered.wait(1)) - self.assertTrue(slot.cancel()) - waiter.join(2) - - self.assertFalse(waiter.is_alive()) - self.assertEqual(len(errors), 1) - self.assertIsInstance(errors[0], TransferCancelledError) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_waits_are_pinned_to_the_started_generation(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-start-fused-only", - thread_mode="multiple", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-start-fused-only", - thread_mode="multiple", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - original_wait_until = TransferSlot._wait_until - - def replace_generation_while_waiter_is_paused(start_and_wait: bool) -> None: - entered = threading.Event() - release = threading.Event() - errors: list[BaseException] = [] - - def pause_before_first_observation(self, *args, **kwargs): - entered.set() - if not release.wait(2): - raise RuntimeError("generation-pinning test timed out") - return original_wait_until(self, *args, **kwargs) - - if start_and_wait: - backend_slot.next_state = WorkState.RUNNING - operation = slot.start_and_wait - else: - backend_slot.next_state = WorkState.COMPLETED - slot.start() - operation = slot.wait - - def run_waiter() -> None: - try: - operation() - except BaseException as error: - errors.append(error) - - with mock.patch.object( - TransferSlot, "_wait_until", pause_before_first_observation - ): - waiter = threading.Thread(target=run_waiter) - waiter.start() - self.assertTrue(entered.wait(1)) - backend_slot._state = WorkState.COMPLETED - self.assertEqual(slot.state, WorkState.COMPLETED) - slot.recycle() - backend_slot.next_state = WorkState.COMPLETED - slot.start() - replacement_generation = slot.generation - release.set() - waiter.join(2) - - self.assertFalse(waiter.is_alive()) - self.assertEqual(len(errors), 1) - self.assertIsInstance(errors[0], BusyError) - self.assertIn("changed while waiting", str(errors[0])) - self.assertEqual(slot.generation, replacement_generation) - slot.recycle() - - replace_generation_while_waiter_is_paused(start_and_wait=False) - replace_generation_while_waiter_is_paused(start_and_wait=True) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_start_and_wait_fallback_remains_spin_oriented(self) -> None: - source = Endpoint("source", backend="test-request-slot-start-fused-only") - destination = Endpoint( - "destination", backend="test-request-slot-start-fused-only" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - backend_slot.poll_states = [WorkState.RUNNING] * 130 + [WorkState.COMPLETED] - - with mock.patch("torch.distributed._transfer._api.time.sleep") as sleep: - self.assertIs(slot.start_and_wait(), slot) - sleep.assert_not_called() - self.assertEqual(backend_slot.poll_count, 131) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_wait_retains_cooperative_backoff(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.next_state = WorkState.RUNNING - slot.start() - - def complete_on_yield(delay: float) -> None: - self.assertEqual(delay, 0) - backend_slot._state = WorkState.COMPLETED - - with mock.patch( - "torch.distributed._transfer._api.time.sleep", - side_effect=complete_on_yield, - ) as sleep: - self.assertIs(slot.wait(), slot) - sleep.assert_called_once_with(0) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_rejects_concurrent_start(self) -> None: - source = Endpoint("source", backend="test-request-slot", thread_mode="multiple") - destination = Endpoint( - "destination", backend="test-request-slot", thread_mode="multiple" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.start_entered = threading.Event() - backend_slot.start_release = threading.Event() - errors: list[BaseException] = [] - - def start_slot() -> None: - try: - slot.start() - except BaseException as error: - errors.append(error) - - first = threading.Thread(target=start_slot) - second = threading.Thread(target=start_slot) - first.start() - self.assertTrue(backend_slot.start_entered.wait(1)) - second.start() - backend_slot.start_release.set() - first.join(2) - second.join(2) - - self.assertFalse(first.is_alive()) - self.assertFalse(second.is_alive()) - self.assertEqual(len(errors), 1) - self.assertIsInstance(errors[0], BusyError) - self.assertEqual(backend_slot.start_count, 1) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_close_is_linearized_with_fused_start(self) -> None: - source = Endpoint( - "source", backend="test-request-slot-fused", thread_mode="multiple" - ) - destination = Endpoint( - "destination", backend="test-request-slot-fused", thread_mode="multiple" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - backend_slot.start_entered = threading.Event() - backend_slot.start_release = threading.Event() - start_errors: list[BaseException] = [] - close_errors: list[BaseException] = [] - - def start_slot() -> None: - try: - slot.start_and_wait() - except BaseException as error: - start_errors.append(error) - - def close_slot() -> None: - try: - slot.close() - except BaseException as error: - close_errors.append(error) - - starter = threading.Thread(target=start_slot) - closer = threading.Thread(target=close_slot) - starter.start() - self.assertTrue(backend_slot.start_entered.wait(1)) - closer.start() - backend_slot.start_release.set() - starter.join(2) - closer.join(2) - - self.assertFalse(starter.is_alive()) - self.assertFalse(closer.is_alive()) - self.assertEqual(start_errors, []) - self.assertEqual(len(close_errors), 1) - self.assertIsInstance(close_errors[0], BusyError) - self.assertTrue(slot.active) - self.assertFalse(slot.closed) - - slot.recycle() - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_interrupted_before_provider_ack_is_recyclable(self) -> None: - source = Endpoint("source", backend="test-request-slot-unacknowledged-start") - destination = Endpoint( - "destination", backend="test-request-slot-unacknowledged-start" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - self.assertIsNone(slot.start()) - self.assertEqual(slot.state, WorkState.FAILED) - self.assertIsInstance(slot.error, KeyboardInterrupt) - with self.assertRaises(BackendFailureError) as context: - slot.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - slot.recycle() - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_normal_return_without_provider_ack_is_recyclable( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-no-ack") - destination = Endpoint("destination", backend="test-request-slot-no-ack") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - self.assertIsNone(slot.start()) - self.assertEqual(slot.state, WorkState.FAILED) - with self.assertRaisesRegex(BackendFailureError, "without acknowledging"): - slot.wait() - slot.recycle() - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_fused_interruption_before_ack_is_recyclable(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused-unacknowledged") - destination = Endpoint( - "destination", backend="test-request-slot-fused-unacknowledged" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - with self.assertRaises(BackendFailureError) as context: - slot.start_and_wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - self.assertEqual(slot.state, WorkState.FAILED) - self.assertTrue(slot.active) - slot.recycle() - self.assertFalse(slot.active) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_single_mode_rejects_every_wrong_thread_operation( - self, - ) -> None: - source = Endpoint( - "source", backend="test-request-slot-fused", thread_mode="single" - ) - destination = Endpoint( - "destination", backend="test-request-slot-fused", thread_mode="single" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - def wrong_thread_error(function) -> BaseException: - errors: list[BaseException] = [] - - def invoke() -> None: - try: - function() - except BaseException as error: - errors.append(error) - - thread = threading.Thread(target=invoke) - thread.start() - thread.join(2) - self.assertFalse(thread.is_alive()) - self.assertEqual(len(errors), 1) - return errors[0] - - for function in ( - slot.start, - slot.start_and_wait, - lambda: slot.state, - slot.test, - slot.recycle, - slot.close, - ): - self.assertIsInstance(wrong_thread_error(function), UnsupportedError) - - slot.start() - slot.recycle() - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_lifetime_blockers_and_recovery_snapshot(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - self.assertEqual(source.open_transfer_slots, (slot,)) - with self.assertRaisesRegex(BusyError, "request slots"): - bound.close() - with self.assertRaisesRegex(BusyError, "transfer slot"): - source.close() - slot.start() - with self.assertRaisesRegex(BusyError, "unrecycled generation"): - slot.close() - slot.recycle() - slot.close() - self.assertEqual(source.open_transfer_slots, ()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_host_wait_runs_on_every_fused_start(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local.region() - local_region._device = "cuda:0" - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - actions: list[str] = [] - synchronize_bind_count = 0 - - class SynchronizeDescriptor: - def __get__(self, instance: object, owner: object) -> object: - nonlocal synchronize_bind_count - synchronize_bind_count += 1 - return instance._synchronize - - class FakeEvent: - device = torch.device("cuda:0") - synchronize = SynchronizeDescriptor() - - def _synchronize(self) -> None: - actions.append("wait") - - with ( - mock.patch.object(torch.cuda, "Event", FakeEvent), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=False - ) as capture_query, - ): - bound = plan.bind( - TransferOp.WRITE, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(), - ) - slot = bound.create_slot() - self.assertEqual(synchronize_bind_count, 1) - for _ in range(2): - slot.start_and_wait() - slot.recycle() - - self.assertEqual(actions, ["wait", "wait"]) - self.assertEqual(synchronize_bind_count, 1) - capture_query.assert_not_called() - slot.close() - self.assertIsNone(slot._ready_event_synchronize) - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_cached_ready_event_failure_precedes_generation(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local.region() - local_region._device = "cuda:0" - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - - class FailingEvent: - device = torch.device("cuda:0") - - def synchronize(self) -> None: - raise RuntimeError("injected CUDA wait failure") - - with mock.patch.object(torch.cuda, "Event", FailingEvent): - bound = plan.bind( - TransferOp.WRITE, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FailingEvent(), - ) - slot = bound.create_slot() - with self.assertRaisesRegex(BackendFailureError, "CUDA readiness"): - slot.start_and_wait() - - self.assertEqual(slot.generation, 0) - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertEqual(source._backend.request_slots[0].start_count, 0) - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_host_wait_stream_reuses_one_event(self) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local.region() - local_region._device = "cuda:0" - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - actions: list[tuple[str, int]] = [] - event_count = 0 - - class FakeEvent: - device = torch.device("cuda:0") - - def __init__(self) -> None: - nonlocal event_count - event_count += 1 - self.event_id = event_count - - def record(self, stream: object) -> None: - self.assert_stream(stream) - actions.append(("record", self.event_id)) - - def assert_stream(self, stream: object) -> None: - if not isinstance(stream, FakeStream): - raise AssertionError("readiness event recorded on wrong stream") - - def synchronize(self) -> None: - actions.append(("wait", self.event_id)) - - class FakeStream: - device = torch.device("cuda:0") - - def is_capturing(self) -> bool: - actions.append(("capture", event_count)) - return False - - with ( - mock.patch.object(torch, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Event", FakeEvent), - ): - bound = plan.bind( - TransferOp.WRITE, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - stream=FakeStream(), - ) - slot = bound.create_slot() - self.assertEqual(event_count, 1) - self.assertEqual(slot.execute(), 1) - self.assertEqual(slot.execute(), 2) - - self.assertEqual( - actions, - [ - ("capture", 1), - ("record", 1), - ("wait", 1), - ("capture", 1), - ("record", 1), - ("wait", 1), - ], - ) - self.assertEqual(source._backend.request_slots[0].start_count, 2) - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_reusable_stream_event_failure_precedes_generation( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-fused") - destination = Endpoint("destination", backend="test-request-slot-fused") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local.region() - local_region._device = "cuda:0" - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - event_count = 0 - wait_count = 0 - - class FakeEvent: - device = torch.device("cuda:0") - - def __init__(self) -> None: - nonlocal event_count - event_count += 1 - - def record(self, _stream: object) -> None: - return None - - def synchronize(self) -> None: - nonlocal wait_count - wait_count += 1 - if wait_count == 1: - raise RuntimeError("injected reusable-event wait failure") - - class FakeStream: - device = torch.device("cuda:0") - - def is_capturing(self) -> bool: - return False - - with ( - mock.patch.object(torch, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Event", FakeEvent), - ): - bound = plan.bind( - TransferOp.WRITE, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - stream=FakeStream(), - ) - slot = bound.create_slot() - with self.assertRaisesRegex(BackendFailureError, "CUDA readiness"): - slot.execute() - self.assertEqual(slot.generation, 0) - self.assertFalse(slot.active) - self.assertEqual(source._backend.request_slots[0].start_count, 0) - self.assertEqual(slot.execute(), 1) - - self.assertEqual(event_count, 1) - self.assertEqual(wait_count, 2) - self.assertEqual(source._backend.request_slots[0].start_count, 1) - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_creation_interruption_rolls_back_or_is_recoverable( - self, - ) -> None: - for backend_name in ( - "test-request-slot-before-adoption-interrupt", - "test-request-slot-create-interrupt", - ): - source = Endpoint("source", backend=backend_name) - destination = Endpoint("destination", backend=backend_name) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region()], remote=[peer.region("buffer")] - ) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaises(KeyboardInterrupt): - bound.create_slot() - self.assertEqual(source.open_transfer_slots, ()) - self.assertFalse(bound._slots) - if backend_name == "test-request-slot-create-interrupt": - self.assertEqual(source._backend.created_slot.close_count, 1) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - source = Endpoint("source", backend="test-request-slot-cleanup-interrupt") - destination = Endpoint( - "destination", backend="test-request-slot-cleanup-interrupt" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - - with self.assertRaisesRegex(BackendFailureError, "cleanup is pending"): - bound.create_slot() - (slot,) = source.open_transfer_slots - with self.assertRaisesRegex(BackendFailureError, "cleanup only"): - slot.start() - slot.close() - self.assertEqual(source.open_transfer_slots, ()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_creation_publication_interrupt_has_no_hidden_blocker( - self, - ) -> None: - class InterruptAfterAdd(set): - def add(self, value) -> None: - super().add(value) - raise KeyboardInterrupt - - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - bound._slots = InterruptAfterAdd(bound._slots) - - with self.assertRaises(KeyboardInterrupt): - bound.create_slot() - self.assertEqual(source.open_transfer_slots, ()) - self.assertFalse(bound._slots) - self.assertEqual(source._backend.request_slots, []) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_close_commit_interruptions_are_retryable(self) -> None: - class InterruptingSet(set): - def __init__(self, values, *, before: bool) -> None: - super().__init__(values) - self.before = before - self.armed = True - - def discard(self, value) -> None: - if self.armed: - self.armed = False - if self.before: - raise KeyboardInterrupt - super().discard(value) - raise KeyboardInterrupt - super().discard(value) - - for boundary in ("endpoint-after", "bound-before", "bound-after"): - with self.subTest(boundary=boundary): - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region()], remote=[peer.region("buffer")] - ) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - if boundary == "endpoint-after": - source._transfer_slots = InterruptingSet( - source._transfer_slots, before=False - ) - else: - bound._slots = InterruptingSet( - bound._slots, before=boundary == "bound-before" - ) - - with self.assertRaises(KeyboardInterrupt): - slot.close() - if boundary == "endpoint-after": - self.assertTrue(slot.closed) - self.assertNotIn(slot, bound._slots) - else: - self.assertTrue(slot.closed) - slot.close() - self.assertTrue(slot.closed) - self.assertIsNone(slot._bound) - self.assertNotIn(slot, source._transfer_slots) - self.assertNotIn(slot, bound._slots) - self.assertGreaterEqual(backend_slot.close_count, 1) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_pre_provider_interrupt_recycles_core_only(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - original_begin = TransferSlot._begin_generation_unlocked - - def interrupt_after_core_publication(self) -> None: - original_begin(self) - raise KeyboardInterrupt - - with mock.patch.object( - TransferSlot, - "_begin_generation_unlocked", - interrupt_after_core_publication, - ): - self.assertIsNone(slot.start()) - - self.assertTrue(slot.active) - self.assertEqual(slot.generation, 1) - self.assertEqual(backend_slot.start_count, 0) - self.assertEqual(slot.state, WorkState.FAILED) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 0) - - backend_slot.next_state = WorkState.COMPLETED - slot.start() - self.assertEqual(slot.generation, 2) - slot.recycle() - self.assertEqual(backend_slot.recycle_count, 1) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_recycle_interrupt_publishes_clean_idle_last(self) -> None: - class InterruptingIdleCommitTransferSlot(TransferSlot): - __slots__ = ("interrupt_idle_commit",) - - def __init__(self, *args, **kwargs) -> None: - self.interrupt_idle_commit = False - super().__init__(*args, **kwargs) - - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - if ( - name == "_active" - and value is False - and getattr(self, "interrupt_idle_commit", False) - ): - object.__setattr__(self, "interrupt_idle_commit", False) - raise KeyboardInterrupt - - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - with mock.patch.object( - _api, "TransferSlot", InterruptingIdleCommitTransferSlot - ): - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - prior_error = RuntimeError("prior generation failed") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = prior_error - slot.start() - self.assertTrue(slot.test()) - self.assertIs(slot.error, prior_error) - - slot.interrupt_idle_commit = True - with self.assertRaises(KeyboardInterrupt): - slot.recycle() - - self.assertFalse(slot.active) - self.assertIsNone(slot.state) - self.assertIsNone(slot.error) - # The public idle commit precedes private cleanup. The retained - # terminal state is an explicit dirty-idle marker that the next begin - # must repair before publishing another generation. - self.assertEqual(slot._state, WorkState.FAILED) - self.assertTrue(slot._backend_generation_active) - self.assertIs(slot._error, prior_error) - - backend_slot.next_state = WorkState.COMPLETED - backend_slot.next_error = None - self.assertEqual(slot.execute(), 2) - self.assertIsNone(slot.error) - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_terminal_error_publication_retries_after_interrupt( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - backend_slot = source._backend.request_slots[0] - - failure = RuntimeError("retryable provider failure") - backend_slot.next_state = WorkState.FAILED - backend_slot.next_error = failure - backend_slot.raise_error_once = KeyboardInterrupt() - slot.start() - with self.assertRaises(KeyboardInterrupt): - _ = slot.state - self.assertEqual(slot._state, WorkState.PENDING) - self.assertEqual(slot.state, WorkState.FAILED) - self.assertIs(slot.error, failure) - slot.recycle() - - backend_slot.next_state = WorkState.CANCELLED - slot.start() - with ( - mock.patch( - "torch.distributed._transfer._api.TransferCancelledError", - side_effect=KeyboardInterrupt, - ), - self.assertRaises(KeyboardInterrupt), - ): - _ = slot.state - self.assertEqual(slot._state, WorkState.PENDING) - self.assertEqual(slot.state, WorkState.CANCELLED) - self.assertIsInstance(slot.error, TransferCancelledError) - slot.recycle() - - backend_slot.next_state = None - slot.start() - with ( - mock.patch( - "torch.distributed._transfer._api.BackendFailureError", - side_effect=KeyboardInterrupt, - ), - self.assertRaises(KeyboardInterrupt), - ): - _ = slot.state - self.assertEqual(slot._state, WorkState.PENDING) - self.assertEqual(slot.state, WorkState.FAILED) - self.assertIsInstance(slot.error, BackendFailureError) - slot.recycle() - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_transfer_slot_completion_event_is_capability_gated_and_generation_local( - self, - ) -> None: - source = Endpoint("source", backend="test-request-slot-no-completion-event") - destination = Endpoint( - "destination", backend="test-request-slot-no-completion-event" - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - slot = bound.create_slot() - - with self.assertRaisesRegex(BusyError, "no active generation"): - _ = slot.completion_event - slot.start() - self.assertIsNone(slot.completion_event) - slot.recycle() - with self.assertRaisesRegex(BusyError, "no active generation"): - _ = slot.completion_event - - slot.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_reference_retains_remote_owner_during_copy(self) -> None: - class WeakBytearray(bytearray): - pass - - source = Endpoint("source", thread_mode="multiple") - destination = Endpoint("destination", thread_mode="multiple") - source_registration = source.register(bytearray(b"data"), name="buffer") - destination_buffer = WeakBytearray(4) - destination_owner = weakref.ref(destination_buffer) - destination_registration = destination.register( - destination_buffer, name="buffer" - ) - del destination_buffer - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[source_registration.region()], - remote=[peer.region("buffer")], - ) - - entered_copy = threading.Event() - continue_copy = threading.Event() - original_copy = _reference._copy_spans - - def blocking_copy(source_spans, destination_spans, *, anchors=()): - entered_copy.set() - self.assertTrue(continue_copy.wait(1)) - self.assertIsNotNone(destination_owner()) - return original_copy(source_spans, destination_spans, anchors=anchors) - - _reference._copy_spans = blocking_copy - result = [] - writer = threading.Thread(target=lambda: result.append(source.write(plan))) - try: - writer.start() - self.assertTrue(entered_copy.wait(1)) - destination_registration.close() - gc.collect() - self.assertIsNotNone(destination_owner()) - continue_copy.set() - writer.join(1) - self.assertFalse(writer.is_alive()) - result[0].wait() - finally: - continue_copy.set() - writer.join(1) - _reference._copy_spans = original_copy - - result[0].close() - plan.close() - peer.close() - source_registration.close() - source.close() - destination.close() - - def test_reference_tensor_and_storage_registration(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_tensor = torch.arange(8, dtype=torch.uint8) - destination_tensor = torch.zeros(8, dtype=torch.uint8) - source_registration = source.register( - source_tensor.untyped_storage(), name="storage" - ) - destination_registration = destination.register( - destination_tensor, name="tensor" - ) - peer = destination.import_peer(bytes(source.export_metadata())) - plan = destination.prepare( - local=[destination_registration.region()], - remote=[peer.region("storage")], - ) - - work = destination.read(plan) - work.wait() - self.assertTrue(torch.equal(destination_tensor, source_tensor)) - - work.close() - plan.close() - peer.close() - destination_registration.close() - source_registration.close() - destination.close() - source.close() - - def test_tensor_rebinding_keeps_registered_storage_alive(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_tensor = torch.arange(8, dtype=torch.uint8) - expected = source_tensor.clone() - destination_tensor = torch.zeros_like(source_tensor) - source_registration = source.register(source_tensor, name="source") - destination_registration = destination.register( - destination_tensor, name="destination" - ) - source_tensor.set_(torch.full_like(source_tensor, 99)) - peer = destination.import_peer(bytes(source.export_metadata())) - plan = destination.prepare( - local=[destination_registration.region()], - remote=[peer.region("source")], - ) - - work = destination.read(plan) - work.wait() - self.assertTrue(torch.equal(destination_tensor, expected)) - - work.close() - plan.close() - peer.close() - destination_registration.close() - source_registration.close() - destination.close() - source.close() - - def test_indexed_plan_reuse(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"00112233") - destination_buffer = bytearray(b"----------") - source_registration = source.register(source_buffer, name="blocks") - destination_registration = destination.register( - destination_buffer, name="blocks" - ) - peer = source.import_peer(bytes(destination.export_metadata())) - local = [source_registration.region(index * 2, 2) for index in range(4)] - remote = [peer.region("blocks", index * 2, 2) for index in range(5)] - plan = source.prepare(local=local, remote=remote, indexed=True) - - first = source.write(plan, local_indices=[3, 1], remote_indices=[0, 2]) - self.assertTrue(first.test()) - self.assertEqual(destination_buffer, bytearray(b"33--11----")) - first.close() - second = source.write(plan, local_indices=[0], remote_indices=[4]) - second.wait() - self.assertEqual(destination_buffer, bytearray(b"33--11--00")) - - second.close() - plan.close() - peer.close() - source_registration.close() - destination_registration.close() - source.close() - destination.close() - - def test_index_selection_snapshots_packed_int32_and_reuses_dynamic_tag( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"a0b0c0d0") - destination_buffer = bytearray(b"--------") - local = source.register(source_buffer, name="source") - remote = destination.register(destination_buffer, name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=2, count=4)], - remote=[peer.region("destination", 0, 1, stride=2, count=4)], - indexed=True, - ) - local_indices = array("i", [3, 1]) - remote_indices = array("i", [0, 2]) - self.assertEqual(local_indices.itemsize, 4) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=local_indices, - remote_indices=remote_indices, - ) - self.assertIsInstance(selection, IndexSelection) - self.assertEqual(selection.local_count, 2) - self.assertEqual(selection.remote_count, 2) - self.assertEqual(selection.selected_nbytes, 2) - backend_selection = selection._backend_selection - self.assertIsInstance(backend_selection, BackendIndexSelection) - self.assertTrue(backend_selection.local_indices.readonly) - self.assertTrue(backend_selection.remote_indices.readonly) - self.assertIsInstance(backend_selection.local_indices.obj, bytes) - self.assertIs( - backend_selection.local_indices.obj, - backend_selection.remote_indices.obj, - ) - with self.assertRaisesRegex(AttributeError, "immutable"): - selection._op = TransferOp.READ - with self.assertRaisesRegex(AttributeError, "immutable"): - backend_selection._destination_indices_unique = False - - # The caller-owned mutable arrays no longer affect the snapshot. - local_indices[:] = array("i", [0, 0]) - remote_indices[:] = array("i", [1, 1]) - work = selection.submit(notification=b"first") - self.assertIn(selection, work._dependencies) - work.wait() - self.assertEqual(destination_buffer, bytearray(b"d---b---")) - self.assertEqual( - [item.payload for item in destination.poll_notifications()], - [b"first"], - ) - work.close() - - destination_buffer[:] = b"--------" - bound = selection.bind() - self.assertIs(bound._selection, selection) - work = bound.submit(notification=b"second") - work.wait() - self.assertEqual(destination_buffer, bytearray(b"d---b---")) - self.assertEqual( - [item.payload for item in destination.poll_notifications()], - [b"second"], - ) - work.close() - bound.close() - - immutable_bound = plan.bind( - TransferOp.WRITE, - local_indices=[3, 1], - remote_indices=[0, 2], - ) - with self.assertRaisesRegex(TypeError, "IndexSelection binding"): - immutable_bound.submit(notification=b"override") - immutable_bound.close() - plan.close() - with self.assertRaises(ClosedError): - selection.submit() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_fused_submit_and_bounded_work_poll(self) -> None: - source = Endpoint("source", backend="test-fused-selection") - destination = Endpoint("destination", backend="test-fused-selection") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - ) - backend = source._backend - self.assertIsInstance(backend, _FusedSelectionReferenceBackend) - - work, terminal = selection.submit_and_test( - notification=b"pending", - max_polls=17, - timeout_ns=23, - ) - self.assertFalse(terminal) - backend_work = backend.last_work - self.assertIsNotNone(backend_work) - assert backend_work is not None - self.assertEqual(backend_work.state_queries, 0) - self.assertEqual( - work.poll_bounded( - max_polls=31, - timeout_ns=37, - timeout_check_interval=7, - ), - WorkState.COMPLETED, - ) - self.assertEqual(backend_work.bounded_calls, [(31, 37, 7)]) - self.assertEqual(backend_work.state_queries, 0) - self.assertEqual(work.state, WorkState.COMPLETED) - work.close() - - backend.initial_state = WorkState.COMPLETED - work, state = selection.submit_and_poll(max_polls=0) - self.assertIs(state, WorkState.COMPLETED) - immediate_work = backend.last_work - assert immediate_work is not None - self.assertTrue(work.test()) - self.assertEqual(immediate_work.state_queries, 0) - self.assertEqual( - backend.fused_calls, - [(17, 23, b"pending"), (0, None, None)], - ) - - work.close() - work, terminal = selection.submit_and_test(max_polls=0) - self.assertTrue(terminal) - self.assertEqual(backend.last_work.state_queries, 0) - - for arguments, error_type in ( - ({"max_polls": True}, TypeError), - ({"max_polls": -1}, ValueError), - ({"timeout_ns": True}, TypeError), - ({"timeout_ns": -1}, ValueError), - ): - with self.assertRaises(error_type): - selection.submit_and_test(**arguments) - with self.assertRaises(TypeError): - work.test_bounded(max_polls=True) - with self.assertRaises(ValueError): - work.test_bounded(max_polls=0) - work.close() - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_plan_one_shot_indices_use_private_snapshot_and_fallback(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"a0b0c0d0") - destination_buffer = bytearray(b"--------") - local = source.register(source_buffer, name="source") - remote = destination.register(destination_buffer, name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=2, count=4)], - remote=[peer.region("destination", 0, 1, stride=2, count=4)], - indexed=True, - ) - local_indices = array("i", [3, 1]) - remote_indices = array("i", [0, 2]) - backend = source._backend - original_submit = backend.submit_prevalidated - captured: list[BackendIndexSelection] = [] - - def capture_submit(*args, **kwargs) -> None: - captured.append(kwargs["selection"]) - local_indices[:] = array("i", [0, 0]) - remote_indices[:] = array("i", [1, 1]) - original_submit(*args, **kwargs) - - backend.submit_prevalidated = capture_submit - work = plan.submit_indices( - TransferOp.WRITE, - local_indices=local_indices, - remote_indices=remote_indices, - notification=b"one-shot", - ) - self.assertEqual(destination_buffer, bytearray(b"d---b---")) - self.assertEqual(len(captured), 1) - backend_selection = captured[0] - self.assertIn(backend_selection, work._dependencies) - self.assertFalse( - any(isinstance(value, IndexSelection) for value in work._dependencies) - ) - self.assertTrue(backend_selection.local_indices.readonly) - self.assertTrue(backend_selection.remote_indices.readonly) - self.assertEqual(list(backend_selection.local_indices), [3, 1]) - self.assertEqual(list(backend_selection.remote_indices), [0, 2]) - self.assertEqual( - [item.payload for item in destination.poll_notifications()], - [b"one-shot"], - ) - work.close() - - destination_buffer[:] = b"--------" - work, terminal = plan.submit_indices_and_test( - TransferOp.WRITE, - local_indices=array("i", [2]), - remote_indices=array("i", [3]), - max_polls=0, - ) - # The source-compatible provider cannot expose its post-derived state - # without a poll, so zero remains conservatively nonterminal. - self.assertFalse(terminal) - self.assertTrue(work.test()) - self.assertEqual(destination_buffer, bytearray(b"------c-")) - work.close() - - destination_buffer[:] = b"--------" - work, state = plan.submit_indices_and_poll( - TransferOp.WRITE, - local_indices=array("i", [1]), - remote_indices=array("i", [1]), - max_polls=1, - timeout_ns=None, - ) - self.assertIs(state, WorkState.COMPLETED) - self.assertEqual(destination_buffer, bytearray(b"--b-----")) - work.close() - - # caller_ready never consumes local device identities. In particular, - # a nonuniform-device catalog must not trigger an O(k) region walk just - # to build facts that this ordering mode discards. - plan._local_uniform_device = None - with mock.patch.object( - _api, - "_select_region_blocks", - side_effect=AssertionError("caller_ready walked local devices"), - ): - work = plan.submit_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - ) - work.close() - - with self.assertRaisesRegex(ValueError, "non-empty"): - plan.submit_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - notification=b"", - ) - with self.assertRaisesRegex(ValueError, "does not accept"): - plan.submit_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - stream=object(), - ) - for method_name in ( - "submit_indices", - "submit_indices_and_poll", - "submit_indices_and_test", - ): - method = getattr(plan, method_name) - with self.subTest(method=method_name, argument="op"): - with self.assertRaisesRegex(ValueError, "op must be one of"): - method( - None, - local_indices=[0], - remote_indices=[0], - ) - with self.subTest(method=method_name, argument="local_indices"): - with self.assertRaisesRegex( - InvalidRegionError, "local_indices must be non-empty" - ): - method( - TransferOp.WRITE, - local_indices=None, - remote_indices=[0], - ) - with self.subTest(method=method_name, argument="remote_indices"): - with self.assertRaisesRegex( - InvalidRegionError, "remote_indices must be non-empty" - ): - method( - TransferOp.WRITE, - local_indices=[0], - remote_indices=None, - ) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_plan_one_shot_indices_enter_fused_provider_boundary(self) -> None: - source = Endpoint("source", backend="test-fused-selection") - destination = Endpoint("destination", backend="test-fused-selection") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - backend = source._backend - self.assertIsInstance(backend, _FusedSelectionReferenceBackend) - backend.initial_state = WorkState.COMPLETED - - work, state = plan.submit_indices_and_poll( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - notification=b"fused-one-shot", - max_polls=0, - timeout_ns=None, - ) - self.assertIs(state, WorkState.COMPLETED) - self.assertEqual(backend.fused_calls, [(0, None, b"fused-one-shot")]) - self.assertTrue( - any( - isinstance(value, BackendIndexSelection) for value in work._dependencies - ) - ) - self.assertFalse( - any(isinstance(value, IndexSelection) for value in work._dependencies) - ) - work.close() - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_fused_selection_immediate_state_retains_epilogue_interrupt(self) -> None: - class InterruptOnceOnDiscard(set): - armed = True - - def discard(self, value) -> None: - super().discard(value) - if self.armed: - self.armed = False - raise KeyboardInterrupt - - source = Endpoint("source", backend="test-fused-selection") - destination = Endpoint("destination", backend="test-fused-selection") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - ) - backend = source._backend - self.assertIsInstance(backend, _FusedSelectionReferenceBackend) - backend.initial_state = WorkState.COMPLETED - plan._submitting = InterruptOnceOnDiscard() - - work, terminal = selection.submit_and_test(max_polls=0) - self.assertTrue(terminal) - self.assertFalse(plan._submitting) - self.assertEqual(source.open_works, (work,)) - self.assertEqual(work.state, WorkState.FAILED) - self.assertIsInstance(work.error, KeyboardInterrupt) - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_fused_selection_pending_state_reconciles_epilogue_interrupt(self) -> None: - class InterruptOnceOnDiscard(set): - armed = True - - def discard(self, value) -> None: - super().discard(value) - if self.armed: - self.armed = False - raise KeyboardInterrupt - - source = Endpoint("source", backend="test-fused-selection") - destination = Endpoint("destination", backend="test-fused-selection") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - ) - backend = source._backend - self.assertIsInstance(backend, _FusedSelectionReferenceBackend) - backend.initial_state = WorkState.PENDING - plan._submitting = InterruptOnceOnDiscard() - - work, terminal = selection.submit_and_test(max_polls=0) - self.assertFalse(terminal) - self.assertFalse(plan._submitting) - self.assertEqual(source.open_works, (work,)) - self.assertTrue(work.test_bounded(max_polls=1, timeout_ns=None)) - self.assertEqual(work.state, WorkState.FAILED) - self.assertIsInstance(work.error, KeyboardInterrupt) - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_accepts_index_scalars_and_rejects_unsafe_values( - self, - ) -> None: - class IndexScalar: - def __init__(self, value: int) -> None: - self.value = value - - def __index__(self) -> int: - return self.value - - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - - selection = plan.select_indices( - "write", - local_indices=[IndexScalar(3), IndexScalar(1)], - remote_indices=[IndexScalar(0), IndexScalar(2)], - ) - work = selection.submit() - work.close() - for local_indices, remote_indices, message in ( - ([True], [0], "integers"), - ([0], [False], "integers"), - ([4], [0], "outside"), - ([0, 1], [2, 2], "overlap"), - ([], [], "non-empty"), - ): - with self.assertRaisesRegex(InvalidRegionError, message): - plan.select_indices( - TransferOp.WRITE, - local_indices=local_indices, - remote_indices=remote_indices, - ) - - # The raw convenience fallback now follows the same operator.index - # contract, without requiring callers to construct a selection. - work = source.write( - plan, - local_indices=[IndexScalar(0)], - remote_indices=[IndexScalar(1)], - ) - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_native_int32_analyzer_matches_fallback(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - block_count = _api._INDEX_SELECTION_TORCH_THRESHOLD + 3 - local = source.register(bytearray(block_count), name="source") - remote = destination.register(bytearray(block_count), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=block_count)], - remote=[peer.region("destination", 0, 1, stride=1, count=block_count)], - indexed=True, - ) - cutoff = _api._INDEX_SELECTION_TORCH_THRESHOLD - signatures = [] - for analyzer in (None, _emulated_compiled_index_analyzer): - with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", range(cutoff + 1)), - remote_indices=array("i", reversed(range(cutoff + 1))), - ) - backend_selection = selection._backend_selection - signatures.append( - ( - bytes(backend_selection.local_indices), - bytes(backend_selection.remote_indices), - selection.local_count, - selection.remote_count, - selection.selected_nbytes, - ) - ) - self.assertEqual(signatures[0], signatures[1]) - - caller_local = array("i", range(cutoff + 1)) - caller_remote = array("i", reversed(range(cutoff + 1))) - - def mutate_after_snapshot(*args, **kwargs): - result = _emulated_compiled_index_analyzer(*args, **kwargs) - caller_local[:] = array("i", [0]) * len(caller_local) - caller_remote[:] = array("i", [0]) * len(caller_remote) - return result - - analyzer = mock.Mock(side_effect=mutate_after_snapshot) - with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=caller_local, - remote_indices=caller_remote, - ) - analyzer.assert_called_once() - self.assertEqual( - list(selection._backend_selection.local_indices), - list(range(cutoff + 1)), - ) - self.assertEqual( - list(selection._backend_selection.remote_indices), - list(reversed(range(cutoff + 1))), - ) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_compiled_analyzer_preserves_error_precedence( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(8), name="source") - remote = destination.register(bytearray(8), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=8)], - remote=[peer.region("destination", 0, 1, stride=1, count=8)], - indexed=True, - ) - cases = ( - (array("i", [-1]), array("i", [8]), "local_indices.*outside"), - (array("i", [0]), array("i", [8]), "remote_indices.*outside"), - ( - array("i", [0, 1, 2]), - array("i", [0, 0]), - "byte totals differ", - ), - (array("i", [0, 1]), array("i", [0, 0]), "overlap"), - ) - for analyzer in (None, _emulated_compiled_index_analyzer): - for local_indices, remote_indices, message in cases: - with ( - self.subTest(analyzer=analyzer, message=message), - mock.patch.object(_api, "_compiled_index_analyzer", analyzer), - self.assertRaisesRegex(InvalidRegionError, message), - ): - plan.select_indices( - TransferOp.WRITE, - local_indices=local_indices, - remote_indices=remote_indices, - ) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_mixed_inputs_preserve_error_order_and_snapshot( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(4), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - - cases = ( - (array("i", [-1]), [True], r"local_indices index -1.*outside"), - (array("i", [4]), [False], r"local_indices index 4.*outside"), - (array("i", [-1]), [4], r"local_indices index -1.*outside"), - (array("i", [-1]), [], r"local_indices index -1.*outside"), - ([True], array("i", [4]), r"local_indices must contain integers"), - ([4], array("i", [-1]), r"local_indices index 4.*outside"), - ([], array("i", [-1]), r"local_indices must be non-empty"), - ) - for local_indices, remote_indices, message in cases: - with ( - self.subTest( - local_type=type(local_indices).__name__, - remote_type=type(remote_indices).__name__, - ), - self.assertRaisesRegex(InvalidRegionError, message), - ): - plan.select_indices( - TransferOp.WRITE, - local_indices=local_indices, - remote_indices=remote_indices, - ) - - caller_local = array("i", [3, 1]) - - class MutatingRemote: - def __iter__(self): - caller_local[:] = array("i", [0, 0]) - return iter((0, 2)) - - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=caller_local, - remote_indices=MutatingRemote(), - ) - self.assertEqual(list(selection._backend_selection.local_indices), [3, 1]) - self.assertEqual(list(selection._backend_selection.remote_indices), [0, 2]) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_compiled_analyzer_exception_is_not_retried( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(4), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - analyzer = mock.Mock(side_effect=RuntimeError("compiled helper failed")) - with ( - mock.patch.object(_api, "_compiled_index_analyzer", analyzer), - self.assertRaisesRegex(RuntimeError, "compiled helper failed"), - ): - plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [0, 1]), - remote_indices=array("i", [2, 3]), - ) - analyzer.assert_called_once() - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_missing_helper_and_third_party_sequences(self) -> None: - class IndexScalar: - def __init__(self, value: int) -> None: - self.value = value - - def __index__(self) -> int: - return self.value - - class ThirdPartyIndices: - def __init__(self, *values: int) -> None: - self.values = values - - def __iter__(self): - return iter(IndexScalar(value) for value in self.values) - - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - - with mock.patch.object(_api, "_compiled_index_analyzer", None): - fallback = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - ) - self.assertEqual(list(fallback._backend_selection.local_indices), [3, 1]) - - analyzer = mock.Mock(side_effect=AssertionError("must not be called")) - with mock.patch.object(_api, "_compiled_index_analyzer", analyzer): - mixed = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [2, 0]), - remote_indices=ThirdPartyIndices(1, 3), - ) - analyzer.assert_not_called() - self.assertEqual(list(mixed._backend_selection.remote_indices), [1, 3]) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_compiled_index_analyzer_defensive_buffer_contract(self) -> None: - analyzer = _api._compiled_index_analyzer - if analyzer is None: - self.skipTest("requires torch._C rebuilt with the transfer helper") - - local = array("i", [3, 1, 2]) - remote = array("i", [0, 0, 2]) - result = analyzer(local, remote, True) - self.assertEqual(result[0], local.tobytes() + remote.tobytes()) - self.assertEqual(result[1:7], (3, 3, 1, 3, 0, 2)) - self.assertTrue(result[7]) - - targeted_destinations = ( - array("i", [5]), - array("i", [5, 5]), - array("i", range(4096)), - array("i", reversed(range(4096))), - array("i", [2, 0, 3, 1]), - array("i", [7, 7, -5, 100]), - array("i", ((index * 2053) % 4096 for index in range(4096))), - array( - "i", - (((index * 37) % 64) * 104729 for index in range(64)), - ), - array("i", [0, 104729, -65537, 209458]), - array("i", [7, 8, 7]), - array("i", [0, 1, 2, 3, 4, 5, 6, 7, -1, 9, 10]), - array("i", [*range(4095), -1]), - array("i", [*reversed(range(1, 4096)), 4096]), - array("i", [-(2**31), 2**31 - 1, 0]), - ) - for destination in targeted_destinations: - expected_duplicate = len(set(destination)) != len(destination) - for destination_is_remote in (False, True): - with self.subTest( - destination_is_remote=destination_is_remote, - destination=list(destination[:8]), - ): - if destination_is_remote: - candidate = analyzer(local, destination, True) - else: - candidate = analyzer(destination, remote, False) - self.assertEqual(candidate[7], expected_duplicate) - destination_bounds = ( - candidate[5:7] if destination_is_remote else candidate[3:5] - ) - self.assertEqual( - destination_bounds, (min(destination), max(destination)) - ) - - invalid_buffers = ( - array("i"), - array("I", [0]), - array("q", [0]), - memoryview(array("i", [0, 1, 2, 3]))[::2], - ) - for invalid in invalid_buffers: - with self.subTest(invalid=invalid), self.assertRaises(ValueError): - analyzer(invalid, array("i", [0]), True) - - def test_index_selection_reuse_token_is_lazy_and_not_redundant(self) -> None: - source = Endpoint("source", backend="test-request-slot") - destination = Endpoint("destination", backend="test-request-slot") - local = source.register(bytearray(b"abcd"), name="source") - remote = destination.register(bytearray(4), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("destination", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [3, 1]), - remote_indices=array("i", [0, 2]), - ) - backend = source._backend - original_submit = backend.submit_prevalidated - observed_tokens: list[object | None] = [] - - def capture_submit(*args, **kwargs) -> None: - observed_tokens.append(kwargs["reuse_token"]) - original_submit(*args, **kwargs) - - backend.submit_prevalidated = capture_submit - self.assertIsNone(selection._reuse_token) - direct = selection.submit() - direct.close() - self.assertEqual(observed_tokens, [None]) - self.assertIsNone(selection._reuse_token) - - first_bound = selection.bind() - self.assertIsNone(first_bound._reuse_token) - slot = first_bound.create_slot() - self.assertIsNone(selection._reuse_token) - slot.close() - - first = first_bound.submit() - token = selection._reuse_token - self.assertIsNotNone(token) - first.close() - second = first_bound.submit() - second.close() - second_bound = selection.bind() - self.assertIsNone(second_bound._reuse_token) - third = second_bound.submit() - third.close() - self.assertEqual(len(observed_tokens), 4) - self.assertIsNone(observed_tokens[0]) - self.assertIs(observed_tokens[1], token) - self.assertIs(observed_tokens[2], token) - self.assertIs(observed_tokens[3], token) - - second_bound.close() - first_bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_index_selection_irregular_layout_keeps_conservative_fallback( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"abc"), name="source") - remote = destination.register(bytearray(3), name="destination") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1), local.region(1, 2)], - remote=[peer.region("destination", 0, 3)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[0, 1], - remote_indices=[0], - ) - self.assertFalse(selection._backend_selection.uniform_pair_layout) - self.assertEqual(selection.selected_nbytes, 3) - work = selection.submit() - work.wait() - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_backend_index_selection_cannot_be_forged(self) -> None: - with self.assertRaisesRegex(TypeError, "created by Core"): - BackendIndexSelection( - TransferOp.WRITE, - object(), - memoryview(array("i", [0])), - memoryview(array("i", [0])), - selected_nbytes=1, - uniform_pair_layout=True, - destination_catalog_nonoverlapping=True, - destination_indices_unique=True, - _token=object(), - ) - - def test_indexed_compact_strided_block_catalog(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"a0b0c0d0") - destination_buffer = bytearray(b"--------") - source_registration = source.register(source_buffer, name="source") - destination_registration = destination.register( - destination_buffer, name="destination" - ) - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[source_registration.region(0, 1, stride=2, count=4)], - remote=[peer.region("destination", 0, 1, stride=2, count=4)], - indexed=True, - ) - - self.assertEqual(plan.local_block_count, 4) - self.assertEqual(plan.remote_block_count, 4) - work = source.write( - plan, - local_indices=[3, 1], - remote_indices=[0, 2], - ) - work.wait() - self.assertEqual(destination_buffer, bytearray(b"d---b---")) - work.close() - with self.assertRaises(InvalidRegionError): - source.write(plan, local_indices=[4], remote_indices=[0]) - - plan.close() - peer.close() - source_registration.close() - destination_registration.close() - source.close() - destination.close() - - def test_indexed_logical_block_sizes_must_match(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - destination_buffer = bytearray(4) - source_registration = source.register(bytearray(b"a0b0"), name="source") - destination_registration = destination.register( - destination_buffer, name="destination" - ) - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[source_registration.region(0, 1, stride=2, count=2)], - remote=[peer.region("destination", 0, 2, stride=2, count=2)], - indexed=True, - ) - - with self.assertRaisesRegex(InvalidRegionError, "byte totals differ"): - source.write(plan, local_indices=[0], remote_indices=[0]) - work = source.write( - plan, - local_indices=[0, 1], - remote_indices=[1], - ) - work.wait() - self.assertEqual(destination_buffer, bytearray(b"\x00\x00ab")) - work.close() - - plan.close() - peer.close() - source_registration.close() - destination_registration.close() - source.close() - destination.close() - - def test_strided_regions(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"a0b0c0d0") - destination_buffer = bytearray(4) - source_registration = source.register(source_buffer, name="source") - destination_registration = destination.register( - destination_buffer, name="destination" - ) - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[source_registration.region(0, 1, stride=2, count=4)], - remote=[peer.region("destination")], - ) - - work = source.write(plan) - work.wait() - self.assertEqual(destination_buffer, bytearray(b"abcd")) - - work.close() - plan.close() - peer.close() - source_registration.close() - destination_registration.close() - source.close() - destination.close() - - def test_notifications_identify_source(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - - with self.assertRaisesRegex(ValueError, "non-empty"): - source.write(plan, notification=b"") - transfer = source.write(plan, notification=b"visible") - transfer.wait() - notices = destination.poll_notifications() - self.assertEqual(len(notices), 1) - self.assertEqual(notices[0].payload, b"visible") - self.assertEqual(notices[0].source_endpoint_id, source.endpoint_id) - self.assertEqual(notices[0].source_incarnation, source.incarnation) - self.assertEqual(notices[0].source, source.identity) - - class OversizedPayload(bytearray): - converted = False - - def __bytes__(self) -> bytes: - self.converted = True - return super().__bytes__() - - oversized = OversizedPayload(b"1234") - with ( - mock.patch("torch.distributed._transfer._api._MAX_NOTIFICATION_BYTES", 3), - self.assertRaisesRegex(ValueError, "size limit"), - ): - source.send_notification(peer, oversized) - self.assertFalse(oversized.converted) - - standalone = source.send_notification(peer, b"standalone") - standalone.wait() - self.assertEqual( - [item.payload for item in destination.poll_notifications()], - [b"standalone"], - ) - - self.assertTrue(source.capabilities.synchronous_notification_send) - self.assertEqual(set(source.open_works), {standalone, transfer}) - sync_payload = b"sync-exact-payload" - self.assertIsNone(source.send_notification_sync(peer, sync_payload)) - self.assertEqual(set(source.open_works), {standalone, transfer}) - sync_notifications = destination.poll_notifications() - self.assertEqual([item.payload for item in sync_notifications], [sync_payload]) - self.assertIs(sync_notifications[0].payload, sync_payload) - - provider_failure = RuntimeError("provider failed after entry") - - def fail_after_entry(*args, **kwargs) -> None: - del args, kwargs - raise provider_failure - - source._send_notification_sync_dispatch = fail_after_entry - with self.assertRaisesRegex( - BackendFailureError, "delivery outcome may be ambiguous" - ) as context: - source.send_notification_sync(peer, b"ambiguous") - self.assertIs(context.exception.__cause__, provider_failure) - - transfer_error = ClosedError("provider is closed") - - def fail_with_transfer_error(*args, **kwargs) -> None: - del args, kwargs - raise transfer_error - - source._send_notification_sync_dispatch = fail_with_transfer_error - with self.assertRaises(ClosedError) as context: - source.send_notification_sync(peer, b"rejected") - self.assertIs(context.exception, transfer_error) - - source._send_notification_sync_dispatch = lambda *args, **kwargs: object() - with self.assertRaisesRegex( - BackendFailureError, "must return None.*outcome may be ambiguous" - ): - source.send_notification_sync(peer, b"invalid-result") - self.assertEqual(set(source.open_works), {standalone, transfer}) - - standalone.close() - transfer.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_attached_only_backend_can_poll_notifications(self) -> None: - endpoint = Endpoint("endpoint", backend="test-attached-only") - destination = Endpoint("destination", backend="test-attached-only") - peer = endpoint.import_peer(bytes(destination.export_metadata())) - self.assertTrue(endpoint.capabilities.attached_notifications) - self.assertFalse(endpoint.capabilities.standalone_notifications) - self.assertFalse(endpoint.capabilities.synchronous_notification_send) - with self.assertRaisesRegex(UnsupportedError, "synchronous notification"): - endpoint.send_notification_sync(peer, b"unsupported") - with self.assertRaisesRegex(UnsupportedError, "synchronous standalone"): - _api.Backend.send_notification_sync( - endpoint._backend, object(), b"unsupported" - ) - self.assertEqual(endpoint.poll_notifications(), []) - peer.close() - endpoint.close() - destination.close() - - def test_multiple_synchronous_notification_sends_overlap_and_pin_peer( - self, - ) -> None: - source = Endpoint("source", thread_mode="multiple") - destination = Endpoint("destination", thread_mode="multiple") - peer = source.import_peer(bytes(destination.export_metadata())) - original_send = source._send_notification_sync_dispatch - self.assertIsNotNone(original_send) - entered_first = threading.Event() - entered_both = threading.Event() - release = threading.Event() - count_lock = threading.Lock() - entered = 0 - errors: list[BaseException] = [] - - def blocking_send(peer_handle, payload) -> None: - nonlocal entered - with count_lock: - entered += 1 - entered_first.set() - if entered == 2: - entered_both.set() - if not release.wait(2): - raise RuntimeError("timed out waiting to release concurrent sends") - assert original_send is not None - original_send(peer_handle, payload) - - source._send_notification_sync_dispatch = blocking_send - - def send(payload: bytes) -> None: - try: - source.send_notification_sync(peer, payload) - except BaseException as error: - errors.append(error) - - threads = [ - threading.Thread(target=send, args=(payload,)) - for payload in (b"first", b"second") - ] - try: - for thread in threads: - thread.start() - self.assertTrue(entered_first.wait(1)) - self.assertTrue( - entered_both.wait(1), - "MULTIPLE synchronous sends were serialized before native entry", - ) - with self.assertRaisesRegex(BusyError, "synchronous notification"): - peer.close() - finally: - release.set() - for thread in threads: - thread.join(2) - - self.assertTrue(all(not thread.is_alive() for thread in threads)) - self.assertEqual(errors, []) - self.assertEqual( - sorted(item.payload for item in destination.poll_notifications()), - [b"first", b"second"], - ) - peer.close() - source.close() - destination.close() - - def test_multiple_synchronous_notification_interrupted_cleanup_cannot_wedge_peer( - self, - ) -> None: - class InterruptSecondEnter: - def __init__(self, lock) -> None: - self.lock = lock - self.enters = 0 - - def __enter__(self): - self.enters += 1 - if self.enters == 2: - raise KeyboardInterrupt - return self.lock.__enter__() - - def __exit__(self, *args): - return self.lock.__exit__(*args) - - source = Endpoint("source", thread_mode="multiple") - destination = Endpoint("destination", thread_mode="multiple") - peer = source.import_peer(bytes(destination.export_metadata())) - original_lock = source._state_lock - source._state_lock = InterruptSecondEnter(original_lock) - retained_error = None - try: - try: - source.send_notification_sync(peer, b"accepted-before-interrupt") - except KeyboardInterrupt as error: - # Retain the traceback, and therefore the completed send frame, - # through peer.close(). Close must use the explicit completion - # bit rather than mistake weakref liveness for native activity. - retained_error = error - else: - self.fail("cleanup-lock interruption did not propagate") - finally: - source._state_lock = original_lock - - gc.collect() - self.assertIsNotNone(retained_error) - self.assertTrue(peer._active_sync_notification_sends) - self.assertTrue( - all( - reference() is None or not reference().active - for reference in peer._active_sync_notification_sends - ) - ) - # close() prunes dead reservations and cannot mistake interrupted - # cleanup for a still-running native call. - peer.close() - self.assertFalse(peer._active_sync_notification_sends) - retained_error = None - source.close() - destination.close() - - def test_multiple_async_notification_send_overlaps_manual_progress(self) -> None: - # This is a backend-entry concurrency test, not a timing benchmark. - source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) - destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) - peer = source.import_peer(bytes(destination.export_metadata())) - backend_work = _PendingWork() - send_entered = threading.Event() - progress_entered = threading.Event() - release_send = threading.Event() - errors: list[BaseException] = [] - returned: list[Work] = [] - - def blocked_send(*args, **kwargs) -> None: - kwargs["adopt_work"](backend_work) - send_entered.set() - if not release_send.wait(2): - raise RuntimeError("test did not release notification send") - - def observe_progress() -> None: - progress_entered.set() - - source._backend.send_notification = blocked_send - source._backend.progress = observe_progress - - def send() -> None: - try: - returned.append(source.send_notification(peer, b"overlap")) - except BaseException as error: - errors.append(error) - - send_thread = threading.Thread(target=send) - send_thread.start() - try: - self.assertTrue(send_entered.wait(1)) - progress_thread = threading.Thread(target=source.progress) - progress_thread.start() - progress_thread.join(1) - self.assertFalse(progress_thread.is_alive()) - self.assertTrue(progress_entered.is_set()) - with self.assertRaises(BusyError): - peer.close() - finally: - release_send.set() - send_thread.join(2) - - self.assertFalse(send_thread.is_alive()) - self.assertEqual(errors, []) - self.assertEqual(len(returned), 1) - backend_work._state = WorkState.COMPLETED - returned[0].close() - peer.close() - source.close() - destination.close() - - def test_multiple_notification_consumer_admission_serializes_and_blocks_close( - self, - ) -> None: - endpoint = Endpoint("consumer", thread_mode=ThreadMode.MULTIPLE) - first_entered = threading.Event() - second_entered = threading.Event() - release_first = threading.Event() - count_lock = threading.Lock() - entries = 0 - errors: list[BaseException] = [] - - def blocked_poll(max_items=None): - del max_items - nonlocal entries - with count_lock: - entries += 1 - if entries == 1: - first_entered.set() - else: - second_entered.set() - if entries == 1 and not release_first.wait(2): - raise RuntimeError("test did not release first notification poll") - return [] - - endpoint._backend.poll_notifications = blocked_poll - - def poll() -> None: - try: - endpoint.poll_notifications() - except BaseException as error: - errors.append(error) - - first = threading.Thread(target=poll) - second = threading.Thread(target=poll) - first.start() - self.assertTrue(first_entered.wait(1)) - second.start() - self.assertFalse(second_entered.wait(0.1)) - - close_errors: list[BaseException] = [] - - def close() -> None: - try: - endpoint.close() - except BaseException as error: - close_errors.append(error) - - close_thread = threading.Thread(target=close) - close_thread.start() - close_thread.join(1) - self.assertFalse(close_thread.is_alive()) - self.assertEqual(len(close_errors), 1) - self.assertIsInstance(close_errors[0], BusyError) - self.assertFalse(endpoint.closed) - - release_first.set() - first.join(2) - second.join(2) - self.assertTrue(second_entered.is_set()) - self.assertTrue(all(not thread.is_alive() for thread in (first, second))) - self.assertEqual(errors, []) - (warmed_reservation,) = endpoint._multiple_backend_call_reservation_pool - endpoint.poll_notifications() - self.assertEqual(len(endpoint._multiple_backend_call_reservation_pool), 1) - self.assertIs( - endpoint._multiple_backend_call_reservation_pool[0], warmed_reservation - ) - endpoint.close() - - def test_notification_batches_are_immutable_bounded_and_source_compatible( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - peer = source.import_peer(bytes(destination.export_metadata())) - first = source.send_notification(peer, b"first") - second_payload = bytearray(b"second") - second = source.send_notification(peer, second_payload) - first.wait() - second.wait() - second_payload[:] = b"mutate" - - with mock.patch.object( - destination._backend, - "poll_notification_batches", - side_effect=AssertionError("zero max_items must not enter the backend"), - ): - self.assertEqual(destination.poll_notification_batches(max_items=0), []) - for invalid, error_type in ( - (True, TypeError), - (-1, ValueError), - (1.0, TypeError), - ): - with self.assertRaises(error_type): - destination.poll_notification_batches(max_items=invalid) # type: ignore[arg-type] - - batches = destination.poll_notification_batches() - self.assertEqual(len(batches), 1) - self.assertIsInstance(batches[0], NotificationBatch) - self.assertEqual(batches[0].payloads, (b"first", b"second")) - self.assertEqual(batches[0].source, source.identity) - - exact_payloads = (b"retained-one", b"retained-two") - retained = BackendNotificationBatch( - exact_payloads, - source_endpoint_id=source.endpoint_id, - source_incarnation=source.incarnation, - ) - self.assertIs(retained.payloads, exact_payloads) - self.assertEqual(retained.total_nbytes, sum(map(len, exact_payloads))) - with ( - mock.patch.object( - destination._backend, - "poll_notification_batches", - return_value=[retained], - ), - mock.patch.object( - sys.modules[Endpoint.__module__], - "NotificationBatch", - side_effect=lambda *args, **kwargs: ( - self.assertFalse(destination._state_lock._is_owned()) - or NotificationBatch(*args, **kwargs) - ), - ), - ): - public = destination.poll_notification_batches() - self.assertIs(public[0].payloads, exact_payloads) - - mutable_payload = bytearray(b"snapshotted") - snapshotted = BackendNotificationBatch([mutable_payload]) - mutable_payload[:] = b"overwritten!" - self.assertEqual(snapshotted.payloads, (b"snapshotted",)) - self.assertEqual(snapshotted.total_nbytes, len(b"snapshotted")) - - class OversizedBytearray(bytearray): - def __bytes__(self) -> bytes: - raise AssertionError("oversized mutable input must not be copied") - - oversized_mutable = OversizedBytearray( - b"x" * (BackendNotificationBatch.MAX_PAYLOAD_BYTES + 1) - ) - with self.assertRaisesRegex(ValueError, "wire size limit"): - BackendNotificationBatch((oversized_mutable,)) - - overrun = BackendNotificationBatch( - (b"one", b"two"), - source_endpoint_id=source.endpoint_id, - source_incarnation=source.incarnation, - ) - with ( - mock.patch.object( - destination._backend, - "poll_notification_batches", - return_value=[overrun], - ), - self.assertRaisesRegex(BackendFailureError, "more notification payloads"), - ): - destination.poll_notification_batches(max_items=1) - - with self.assertRaisesRegex(ValueError, "at least one payload"): - BackendNotificationBatch(()) - with self.assertRaisesRegex(ValueError, "specified together"): - NotificationBatch((b"payload",), source_endpoint_id="incomplete") - - first.close() - second.close() - peer.close() - source.close() - destination.close() - - def test_notification_send_uses_work_adoption_and_direct_dispatch(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - peer = source.import_peer(bytes(destination.export_metadata())) - backend_send = source._backend.send_notification - observed_adoption = None - - def capture_send(*args, **kwargs) -> None: - nonlocal observed_adoption - observed_adoption = kwargs["adopt_work"] - backend_send(*args, **kwargs) - - source._backend.send_notification = capture_send - with mock.patch.object( - source, - "_backend_call", - side_effect=AssertionError( - "standalone notification must use its direct backend helper" - ), - ): - work = source.send_notification(peer, b"direct") - - self.assertIs(observed_adoption.__self__, work) - self.assertIs(observed_adoption.__func__, Work._adopt_backend_work) - work.wait() - work.close() - peer.close() - source.close() - destination.close() - - def test_notification_interruption_after_adoption_retains_peer(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - peer = source.import_peer(bytes(destination.export_metadata())) - backend_work = _PendingWork() - - def adopt_then_interrupt(*args, **kwargs) -> None: - kwargs["adopt_work"](backend_work) - raise KeyboardInterrupt - - source._backend.send_notification = adopt_then_interrupt - recovery_token = object() - work = source.send_notification( - peer, - b"message", - recovery_token=recovery_token, - ) - - self.assertEqual(source.open_works, (work,)) - self.assertIs(work.recovery_token, recovery_token) - with self.assertRaises(BusyError): - peer.close() - backend_work._state = WorkState.COMPLETED - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - - work.close() - self.assertIsNone(work.recovery_token) - peer.close() - source.close() - destination.close() - - def test_lifetime_order_is_enforced(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - - with self.assertRaises(BusyError): - local.close() - with self.assertRaises(BusyError): - peer.close() - with self.assertRaises(BusyError): - plan.close() - with self.assertRaises(BusyError): - source.close() - - work.close() - self.assertIsNone(work._plan) - work.close() - plan.close() - self.assertFalse(plan._local) - self.assertFalse(plan._remote) - plan.close() - peer.close() - self.assertIsNone(peer._handle) - peer.close() - local.close() - self.assertIsNone(local._memory.owner) - local.close() - remote.close() - source.close() - source.close() - destination.close() - - def test_plan_cannot_close_during_submit_bookkeeping(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - submit_entered = threading.Event() - allow_submit = threading.Event() - close_done = threading.Event() - close_errors: list[BaseException] = [] - original_submit = source._backend.submit - - def blocked_submit(*args, **kwargs): - submit_entered.set() - allow_submit.wait() - return original_submit(*args, **kwargs) - - source._backend.submit = blocked_submit - work_items: list[Work] = [] - submit_thread = threading.Thread( - target=lambda: work_items.append(source.write(plan)) - ) - - def close_plan() -> None: - try: - plan.close() - except BaseException as error: - close_errors.append(error) - finally: - close_done.set() - - submit_thread.start() - self.assertTrue(submit_entered.wait(1)) - close_thread = threading.Thread(target=close_plan) - close_thread.start() - self.assertTrue(close_done.wait(1)) - self.assertEqual(len(close_errors), 1) - self.assertIsInstance(close_errors[0], BusyError) - self.assertFalse(plan.closed) - allow_submit.set() - submit_thread.join() - close_thread.join() - self.assertEqual(len(work_items), 1) - self.assertFalse(plan.closed) - - work_items[0].close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_overlapping_transfer_is_request_local_failure(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - shared = bytearray(b"data") - local = source.register(shared, name="local") - remote = destination.register(shared, name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - - work = source.write(plan) - self.assertEqual(work.state, WorkState.FAILED) - self.assertIsInstance(work.error, InvalidRegionError) - with self.assertRaises(InvalidRegionError): - work.wait() - - work.close() - self.assertIsNone(work._error) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_indexed_destination_must_not_overlap(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - destination_buffer = bytearray(b"---") - local = source.register(bytearray(b"ab"), name="buffer") - remote = destination.register(destination_buffer, name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=2)], - remote=[ - peer.region("buffer", 0, 1, stride=2, count=2), - peer.region("buffer", 1, 1, stride=1, count=2), - ], - indexed=True, - ) - - work = source.write( - plan, - local_indices=[0, 1], - remote_indices=[0, 2], - ) - work.wait() - self.assertEqual(destination_buffer, bytearray(b"ab-")) - work.close() - - with self.assertRaises(InvalidRegionError): - source.write( - plan, - local_indices=[0, 1], - remote_indices=[1, 3], - ) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_overlapping_registrations_are_rejected(self) -> None: - endpoint = Endpoint("endpoint") - storage = bytearray(b"abcd") - registration = endpoint.register(storage, name="first") - with self.assertRaisesRegex( - InvalidRegionError, "overlaps existing registration" - ): - endpoint.register(storage, name="alias") - registration.close() - endpoint.close() - - def test_reference_deregister_is_idempotent_after_success(self) -> None: - endpoint = Endpoint("endpoint") - registration = endpoint.register(bytearray(4), name="buffer") - handle = registration._handle - - endpoint._backend.deregister(handle) - endpoint._backend.deregister(handle) - registration.close() - endpoint.close() - - def test_one_shot_plan_is_released_with_work(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - - work = source.write(local=[local.region()], remote=[peer.region("buffer")]) - work.wait() - work.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_submit_reservation_is_released_on_base_exception(self) -> None: - source = Endpoint("source", backend="test-interrupt") - destination = Endpoint("destination", backend="test-interrupt") - local = source.register(bytearray(4), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - - with self.assertRaises(KeyboardInterrupt): - source.write(plan) - - self.assertFalse(plan._submitting) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_submit_interruption_after_adoption_retains_lifetimes(self) -> None: - source = Endpoint("source", backend="test-adopt-then-interrupt") - destination = Endpoint("destination", backend="test-adopt-then-interrupt") - local = source.register(bytearray(4), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - - recovery_token = object() - work = source.write(plan, recovery_token=recovery_token) - - self.assertFalse(plan._submitting) - self.assertIs(work.recovery_token, recovery_token) - self.assertIn(work, plan._works) - self.assertIn(work, source._works) - with self.assertRaises(BusyError): - plan.close() - work._backend_work._state = WorkState.COMPLETED - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - - work.close() - self.assertIsNone(work.recovery_token) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_submit_without_work_adoption_is_rolled_back(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(4), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - source._backend.submit = lambda *args, **kwargs: None - - with self.assertRaisesRegex(BackendFailureError, "without adopting"): - source.write(plan) - - self.assertFalse(plan._submitting) - self.assertFalse(plan._works) - self.assertFalse(source._works) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_submit_epilogue_interruption_is_recoverable(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(4), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - - class InterruptOnceOnDiscard(set): - armed = True - - def discard(self, value) -> None: - super().discard(value) - if self.armed: - self.armed = False - raise KeyboardInterrupt - - plan._submitting = InterruptOnceOnDiscard() - work = source.write(plan) - - self.assertFalse(plan._submitting) - self.assertEqual(source.open_works, (work,)) - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIsInstance(context.exception.__cause__, KeyboardInterrupt) - work.close() - self.assertEqual(source.open_works, ()) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_recovery_token_covers_every_work_producing_surface(self) -> None: - class IdentityOnlyToken: - def __hash__(self): - raise AssertionError("recovery token must not be hashed") - - def __eq__(self, other): - raise AssertionError("recovery token must not be compared") +class TestTransfer(unittest.TestCase): + def setUp(self): + self.backend = MemoryBackend() + self.a, self.b = Endpoint(self.backend), Endpoint(MemoryBackend()) + self.src = torch.arange(8, dtype=torch.uint8) + self.dst = torch.zeros_like(self.src) + self.reg = self.a.register_tensor(self.src) + self.b.register_tensor(self.dst) + self.peer = self.a.connect(self.b.metadata()) + self.local = self.a.prepare([(self.src.data_ptr(), 4, 0, 4, 2)], kind="host") + self.remote = self.a.prepare( + [(self.dst.data_ptr(), 4, 0, 4, 2)], kind="host", peer=self.peer + ) - def __bool__(self): - raise AssertionError("recovery token must not be coerced") + def tearDown(self): + self.a.close() + self.b.close() - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"abcd"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("remote", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=array("i", [0]), - remote_indices=array("i", [0]), - ) - bound = selection.bind() - calls = ( - lambda token: plan.submit_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - ), - lambda token: plan.submit_indices_and_test( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - max_polls=0, - ), - lambda token: plan.submit_indices_and_poll( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - max_polls=0, - ), - lambda token: selection.submit(recovery_token=token), - lambda token: selection.submit_and_test(recovery_token=token, max_polls=0), - lambda token: selection.submit_and_poll(recovery_token=token, max_polls=0), - lambda token: bound.submit(recovery_token=token), - lambda token: source.read( - plan, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - ), - lambda token: source.write( - plan, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - ), - lambda token: source.send_notification( - peer, b"recovery", recovery_token=token - ), + def submit(self, operation="WRITE", **kwargs): + return self.a.submit( + operation, self.local, [1, 0], self.remote, [0, 1], **kwargs ) - self.assertEqual(len(calls), 10) - for call in calls: - token = IdentityOnlyToken() - result = call(token) - work = result[0] if isinstance(result, tuple) else result - self.assertIs(work.recovery_token, token) - matches = [ - candidate - for candidate in source.open_works - if candidate.recovery_token is token - ] - self.assertEqual(matches, [work]) + def test_indexed_write_and_reuse(self): + for _ in range(2): + work = self.submit(notification=b"bucket-7") work.wait() + self.assertEqual(self.dst.tolist(), [4, 5, 6, 7, 0, 1, 2, 3]) + self.assertEqual(self.b.notifications(), {self.backend.name: [b"bucket-7"]}) + self.assertEqual(self.b.notifications(), {}) work.close() - self.assertIsNone(work.recovery_token) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_recovery_token_is_strong_and_close_commit_is_retryable(self) -> None: - class Token: - pass - - class InterruptOnceOnDiscard(set): - armed = True - - def discard(self, value) -> None: - super().discard(value) - if self.armed: - self.armed = False - raise KeyboardInterrupt("unlink interrupted") - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - token = Token() - token_ref = weakref.ref(token) - work = source.write(plan, recovery_token=token) - del token - gc.collect() - self.assertIs(token_ref(), work.recovery_token) - work.wait() - - source._works = InterruptOnceOnDiscard(source._works) - with self.assertRaisesRegex(KeyboardInterrupt, "unlink interrupted"): + def test_read_and_standalone_notification(self): + self.dst.fill_(9) + work = self.submit("READ") + work.wait() + self.assertEqual(self.src.tolist(), [9] * 8) + work.close() + self.a.notify(self.peer, b"ack") + self.assertEqual(self.b.notifications(), {self.backend.name: [b"ack"]}) + + def test_close_in_dependency_order(self): + work = self.submit() + for close in ( + self.reg.close, + self.local.close, + lambda: self.a.disconnect(self.peer), + ): + with self.assertRaises(RuntimeError): + close() + work.close() + self.a.close() + self.a.close() + with self.assertRaises(RuntimeError): + self.a.register_tensor(self.src) + + def test_timeout_and_release_failure_retain_resources(self): + self.backend.post = Mock(return_value="pending") + self.backend.poll = Mock(return_value="pending") + work = self.submit() + with self.assertRaises(TimeoutError): + work.wait(timeout=0) + with self.assertRaises(RuntimeError): work.close() - self.assertTrue(work.closed) - self.assertIs(token_ref(), work.recovery_token) - - work.close() - self.assertIsNone(work.recovery_token) - gc.collect() - self.assertIsNone(token_ref()) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_recovery_token_clears_on_every_preadoption_rollback(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"abcd"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("remote", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - ) - bound = plan.bind( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - ) - captured = [] - - def reject_adoption(*args, **kwargs) -> None: - captured.append(kwargs["adopt_work"].__self__) - - cases = ( - ( - "submit", - lambda token: source.write( - plan, - local_indices=[0], - remote_indices=[0], - recovery_token=token, - ), - ), - ("submit", lambda token: bound.submit(recovery_token=token)), - ( - "submit_prevalidated", - lambda token: selection.submit(recovery_token=token), - ), - ( - "send_notification", - lambda token: source.send_notification( - peer, b"rollback", recovery_token=token - ), - ), - ) - for backend_method, submit in cases: - token = object() - with ( - mock.patch.object( - source._backend, - backend_method, - side_effect=reject_adoption, - ), - self.assertRaisesRegex(BackendFailureError, "without adopting"), - ): - submit(token) - self.assertIsNone(captured[-1].recovery_token) - self.assertEqual(source.open_works, ()) - - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_one_shot_validation_failure_does_not_leak_plan(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - - with self.assertRaises(InvalidRegionError): - source.write( - local=[local.region()], - remote=[peer.region("buffer")], - local_indices=[0], - remote_indices=[0], - ) - self.assertFalse(source._plans) - - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_one_shot_plan_release_can_be_retried(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - work = source.write(local=[local.region()], remote=[peer.region("buffer")]) - original_release = source._backend.release_plan - fail_once = True - - def release_plan(plan) -> None: - nonlocal fail_once - if fail_once: - fail_once = False - raise RuntimeError("injected release failure") - original_release(plan) - - source._backend.release_plan = release_plan - with self.assertRaises(BackendFailureError): + self.backend.poll.return_value = "done" + self.backend.release_work = Mock(side_effect=RuntimeError("retry")) + with self.assertRaises(RuntimeError): work.close() - self.assertFalse(work.closed) - self.assertTrue(source._plans) - work.close() - self.assertFalse(source._plans) - - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_failed_one_shot_cleanup_is_retried_by_child_close(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - original_release = source._backend.release_plan - fail_once = True - - def release_plan(plan) -> None: - nonlocal fail_once - if fail_once: - fail_once = False - raise RuntimeError("injected cleanup failure") - original_release(plan) - - source._backend.release_plan = release_plan - with self.assertRaises(InvalidRegionError): - source.write( - local=[local.region()], - remote=[peer.region("buffer")], - local_indices=[0], - remote_indices=[0], - ) - self.assertTrue(source._orphaned_plans) - peer.close() - self.assertFalse(source._orphaned_plans) - - local.close() - remote.close() - source.close() - destination.close() - - def test_rejects_invalid_regions_and_implicit_copies(self) -> None: - endpoint = Endpoint("endpoint") - registration = endpoint.register(bytearray(8), name="buffer") - with self.assertRaises(InvalidRegionError): - registration.region(7, 2) - with self.assertRaises(InvalidRegionError): - registration.region(0, 1, stride=2) - with self.assertRaises(UnsupportedError): - endpoint.register(RawSpan(1, 8, "cpu", "host", bytearray(8)), name="raw") - - tensor = torch.empty((2, 2)).t() - with self.assertRaises(InvalidRegionError): - endpoint.register(tensor, name="noncontiguous") - with self.assertRaises(InvalidRegionError): - endpoint.register(torch.empty(0), name="empty") - - registration.close() - endpoint.close() - - def test_resized_storage_is_rejected_before_use(self) -> None: - endpoint = Endpoint("endpoint") - tensor = torch.arange(8, dtype=torch.uint8) - storage = tensor.untyped_storage() - registration = endpoint.register(tensor, name="tensor") - storage.resize_(16) - with self.assertRaises(InvalidRegionError): - registration.region() - with self.assertRaises(InvalidRegionError): - endpoint.export_metadata() - registration.close() - endpoint.close() + with self.assertRaises(RuntimeError): + self.local.close() + self.backend.release_work.side_effect = None + work.close() + + def test_post_exception_retains_handle_for_endpoint_cleanup(self): + self.backend.post = Mock(side_effect=RuntimeError("post uncertain")) + with self.assertRaises(RuntimeError): + self.submit() + with self.assertRaises(RuntimeError): + self.reg.close() + self.a.close() # reference provider reports terminal completion + + def test_tensor_owner_is_retained(self): + owner = weakref.ref(self.src) + del self.src + gc.collect() + self.assertIsNotNone(owner()) + self.a.close() + gc.collect() + self.assertIsNone(owner()) - def test_cuda_ordering_modes_are_explicit(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) + def test_host_wait_and_invalid_ordering(self): + event = Mock() + self.submit(ordering="host_wait", ready_event=event).close() + event.synchronize.assert_called_once_with() + for kwargs in ( + {"ordering": "host_wait"}, + {"ordering": "unknown"}, + {"ready_event": event}, + ): + with self.assertRaises(ValueError): + self.submit(**kwargs) - self.assertEqual( - source.capabilities.cuda_ordering_modes, - frozenset({CudaOrderingMode.CALLER_READY}), - ) - self.assertTrue(source.capabilities.attached_notifications) - self.assertTrue(source.capabilities.standalone_notifications) - self.assertEqual( - source.capabilities.standalone_notification_completion, - NotificationCompletion.REMOTE_DELIVERED, - ) + def test_invalid_catalogs_indices_and_registration(self): with self.assertRaises(ValueError): - source.write(plan, stream=object()) + self.a.prepare([(self.src.data_ptr() + 8, 4, 0)], kind="host") with self.assertRaises(ValueError): - source.write(plan, cuda_ordering="host_wait") - with self.assertRaises(UnsupportedError): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=object(), - ) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_host_wait_is_core_owned_and_runs_outside_backend_lock(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - source_buffer = bytearray(b"old!") - destination_buffer = bytearray(4) - local_registration = source.register(source_buffer, name="buffer") - remote_registration = destination.register(destination_buffer, name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local_registration.region() - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - local_region._device = "cuda:0" - source._capabilities = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - strided_regions=True, - manual_progress=True, - background_progress=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - standalone_notification_completion=( - NotificationCompletion.REMOTE_DELIVERED - ), - ) - - progressed = threading.Event() - test_case = self - - class FakeEvent: - device = torch.device("cuda:0") - - def synchronize(self) -> None: - progress_thread = threading.Thread( - target=lambda: (source.progress(), progressed.set()) - ) - progress_thread.start() - if not progressed.wait(1): - test_case.fail("CUDA wait ran while the provider lock was held") - progress_thread.join(1) - source_buffer[:] = b"new!" - - with ( - mock.patch.object(torch.cuda, "Event", FakeEvent), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=False - ), - ): - work = source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(), - ) - self.assertEqual(destination_buffer, b"new!") - - work.close() - plan.close() - peer.close() - local_registration.close() - remote_registration.close() - source.close() - destination.close() - - def test_host_wait_stream_records_once_and_wait_failure_rolls_back(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local_registration = source.register(bytearray(b"data"), name="buffer") - remote_registration = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local_registration.region() - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - local_region._device = "cuda:1" - source._capabilities = Capabilities( - read=True, - write=True, - manual_progress=True, - background_progress=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - ) - actions: list[str] = [] - - class FakeEvent: - device = torch.device("cuda:1") - - def synchronize(self) -> None: - actions.append("wait") - - class FakeStream: - device = torch.device("cuda:1") - - def record_event(self) -> FakeEvent: - actions.append("record") - return FakeEvent() - - class FakeStreamContext: - def __enter__(self) -> None: - actions.append("capture-check-enter") - - def __exit__(self, *args: object) -> None: - actions.append("capture-check-exit") - - with ( - mock.patch.object(torch, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Event", FakeEvent), - mock.patch.object(torch.cuda, "stream", return_value=FakeStreamContext()), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=False - ), - ): - work = source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - stream=FakeStream(), - ) - self.assertEqual( - actions, - ["capture-check-enter", "capture-check-exit", "record", "wait"], - ) - work.close() - - actions.clear() - with ( - mock.patch.object(torch, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Stream", FakeStream), - mock.patch.object(torch.cuda, "Event", FakeEvent), - mock.patch.object(torch.cuda, "stream", return_value=FakeStreamContext()), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=True - ), - self.assertRaisesRegex(UnsupportedError, "Graph capture"), - ): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - stream=FakeStream(), - ) - self.assertEqual(actions, ["capture-check-enter", "capture-check-exit"]) - self.assertFalse(plan._submitting) - - class FailingEvent(FakeEvent): - def synchronize(self) -> None: - raise RuntimeError("wait failed") - - with ( - mock.patch.object(torch.cuda, "Event", FailingEvent), - mock.patch.object( - torch.cuda, "is_current_stream_capturing", return_value=False - ), - self.assertRaisesRegex(BackendFailureError, "CUDA readiness"), - ): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FailingEvent(), - ) - self.assertFalse(plan._submitting) - - plan.close() - peer.close() - local_registration.close() - remote_registration.close() - source.close() - destination.close() - - def test_host_wait_rejects_wrong_memory_token_and_skips_event_capture_query( - self, - ) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local_registration = source.register(bytearray(b"data"), name="buffer") - remote_registration = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - local_region = local_registration.region() - plan = source.prepare(local=[local_region], remote=[peer.region("buffer")]) - source._capabilities = Capabilities( - read=True, - write=True, - manual_progress=True, - background_progress=True, - cuda_ordering_modes=frozenset( - {CudaOrderingMode.CALLER_READY, CudaOrderingMode.HOST_WAIT} - ), - ) - - class FakeEvent: - def __init__(self, device: object) -> None: - self.device = device - - def synchronize(self) -> None: - pass - - with mock.patch.object(torch.cuda, "Event", FakeEvent): - with self.assertRaisesRegex(InvalidRegionError, "selected local"): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(torch.device("cuda:0")), - ) - local_region._device = "cuda:0" - with self.assertRaisesRegex(InvalidRegionError, "recorded CUDA"): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(None), - ) - with self.assertRaisesRegex(InvalidRegionError, "does not match"): - source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(torch.device("cuda:1")), - ) - with mock.patch.object( - torch.cuda, - "is_current_stream_capturing", - side_effect=AssertionError("ready-event path queried current stream"), - ) as capture_query: - work = source.write( - plan, - cuda_ordering=CudaOrderingMode.HOST_WAIT, - ready_event=FakeEvent(torch.device("cuda:0")), - ) - capture_query.assert_not_called() - work.close() - self.assertFalse(plan._submitting) - - plan.close() - peer.close() - local_registration.close() - remote_registration.close() - source.close() - destination.close() - - def test_expected_peer_identity_is_checked_before_import(self) -> None: - local = Endpoint("local") - remote = Endpoint("remote", endpoint_id="remote-id") - metadata = bytes(remote.export_metadata()) - - with self.assertRaisesRegex(InvalidMetadataError, "name"): - local.import_peer(metadata, expected_name="other") - self.assertEqual(local._peers, {}) - - peer = local.import_peer( - metadata, - expected_identity=EndpointIdentity("remote-id", remote.incarnation), - expected_name="remote", - ) - peer.close() - local.close() - remote.close() - - def test_work_timeout_and_best_effort_cancel(self) -> None: - source = Endpoint("source", backend="test-pending") - destination = Endpoint("destination", backend="test-pending") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - with self.assertRaises(TransferTimeoutError): - work.wait(0) + self.a.submit("WRITE", self.remote, [0], self.local, [0]) with self.assertRaises(ValueError): - work.wait(float("nan")) - self.assertTrue(work.cancel()) - self.assertEqual(work.state, WorkState.CANCELLED) - with self.assertRaises(TransferCancelledError): - work.wait() - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_zero_timeout_bounded_work_fallback_observes_once(self) -> None: - backend_work = _CountingPendingWork() - self.assertEqual( - backend_work.poll_bounded( - max_polls=64, - timeout_ns=0, - timeout_check_interval=32, - ), - WorkState.PENDING, - ) - self.assertEqual(backend_work.state_queries, 1) - - for max_polls, timeout_ns, timeout_check_interval in ( - (1 << 63, None, 32), - (1, 1 << 63, 32), - (1, None, 1 << 63), - ): - with self.assertRaises(ValueError): - backend_work.poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - - source = Endpoint("source", backend="test-counting-pending") - destination = Endpoint("destination", backend="test-counting-pending") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - self.assertFalse( - work.test_bounded( - max_polls=64, - timeout_ns=0, - timeout_check_interval=32, - ) - ) - provider_work = source._backend.last_work - self.assertIsNotNone(provider_work) - assert provider_work is not None - self.assertEqual(provider_work.state_queries, 1) - for arguments in ( - {"max_polls": 1 << 63}, - {"timeout_ns": 1 << 63}, - {"timeout_check_interval": 1 << 63}, - ): - with self.assertRaises(ValueError): - work.test_bounded(**arguments) - - self.assertTrue(work.cancel()) - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_work_failure_diagnostic_publication_retries_after_interrupt( - self, - ) -> None: - source = Endpoint("source", backend="test-interrupting-failure") - destination = Endpoint("destination", backend="test-interrupting-failure") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - provider_work = source._backend.last_work - self.assertIsNotNone(provider_work) - assert provider_work is not None - - with self.assertRaises(KeyboardInterrupt): - _ = work.state - self.assertEqual(work._state, WorkState.PENDING) - self.assertIsNone(work._error) - self.assertEqual(provider_work.error_queries, 1) - - self.assertEqual(work.state, WorkState.FAILED) - self.assertIs(work.error, provider_work.failure) - with self.assertRaises(BackendFailureError) as context: - work.wait() - self.assertIs(context.exception.__cause__, provider_work.failure) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_lifecycle_constructors_cannot_be_forged(self) -> None: - endpoint = Endpoint("endpoint") - with self.assertRaises(TypeError): - Work(endpoint, _PendingWork()) - endpoint.close() - - def test_background_only_backend_closes_without_manual_progress(self) -> None: - endpoint = Endpoint( - "endpoint", - backend="test-background-only", - progress_mode="background", - ) - with self.assertRaises(UnsupportedError): - endpoint.progress() - endpoint.close() - - def test_backend_error_must_be_an_exception(self) -> None: - source = Endpoint("source", backend="test-bad-error") - destination = Endpoint("destination", backend="test-bad-error") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - - self.assertEqual(work.state, WorkState.FAILED) - self.assertIsInstance(work.error, BackendFailureError) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_backend_work_state_must_not_regress(self) -> None: - source = Endpoint("source", backend="test-regressing") - destination = Endpoint("destination", backend="test-regressing") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - - self.assertEqual(work.state, WorkState.RUNNING) - with self.assertRaisesRegex(BackendFailureError, "regressed"): - _ = work.state - self.assertEqual(work._state, WorkState.RUNNING) - self.assertTrue(work.cancel()) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_state_query_failure_does_not_fake_terminal_completion(self) -> None: - source = Endpoint("source", backend="test-transient-state-error") - destination = Endpoint("destination", backend="test-transient-state-error") - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - - with self.assertRaises(BackendFailureError): - work.test() - self.assertEqual(work._state, WorkState.PENDING) - with self.assertRaises(BusyError): - work.close() - self.assertFalse(plan.closed) - self.assertFalse(local.closed) - self.assertTrue(work.cancel()) - - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_endpoint_id_detects_reincarnation(self) -> None: - local = Endpoint("local") - first = Endpoint("remote", endpoint_id="stable-remote") - old_peer = local.import_peer(bytes(first.export_metadata())) - first.close() - second = Endpoint("remote", endpoint_id="stable-remote") - - with self.assertRaises(BusyError): - local.import_peer(bytes(second.export_metadata())) - old_peer.close() - new_peer = local.import_peer(bytes(second.export_metadata())) - self.assertNotEqual(old_peer.incarnation, new_peer.incarnation) - - new_peer.close() - local.close() - second.close() - - def test_reference_selective_metadata_is_a_registration_grant(self) -> None: - remote = Endpoint("remote") - first = remote.register(bytearray(4), name="first") - second = remote.register(bytearray(4), name="second") - full = remote.export_metadata() - subset = remote.export_metadata([first]) - local = Endpoint("local") - peer = local.import_peer(bytes(subset)) - with self.assertRaises(InvalidRegionError): - peer.region("second") - - second_metadata = next( - item for item in full.registrations if item.name == "second" - ) - spliced = EndpointMetadata( - wire_major=subset.wire_major, - wire_minor=subset.wire_minor, - required_features=subset.required_features, - backend=subset.backend, - endpoint_id=subset.endpoint_id, - name=subset.name, - incarnation=subset.incarnation, - backend_payload=subset.backend_payload, - registrations=(second_metadata,), - ) - other_local = Endpoint("other-local") - with self.assertRaises(InvalidMetadataError): - other_local.import_peer(bytes(spliced)) - - peer.close() - local.close() - other_local.close() - first.close() - second.close() - remote.close() - - def test_reference_rejects_renamed_registration_metadata(self) -> None: - remote = Endpoint("remote") - remote_registration = remote.register(bytearray(4), name="original") - metadata = remote.export_metadata() - original = metadata.registrations[0] - renamed = EndpointMetadata( - wire_major=metadata.wire_major, - wire_minor=metadata.wire_minor, - required_features=metadata.required_features, - backend=metadata.backend, - endpoint_id=metadata.endpoint_id, - name=metadata.name, - incarnation=metadata.incarnation, - backend_payload=metadata.backend_payload, - registrations=( - RegistrationMetadata( - registration_id=original.registration_id, - name="forged", - nbytes=original.nbytes, - device=original.device, - memory_type=original.memory_type, - backend_descriptor=original.backend_descriptor, - ), - ), - ) - local = Endpoint("local") - local_registration = local.register(bytearray(4), name="local") - peer = local.import_peer(renamed) - - with self.assertRaises(InvalidMetadataError): - local.prepare( - local=[local_registration.region()], - remote=[peer.region("forged")], - ) - - peer.close() - local_registration.close() - local.close() - remote_registration.close() - remote.close() - - def test_single_thread_mode_rejects_other_thread(self) -> None: - endpoint = Endpoint("endpoint", thread_mode="single") - errors: list[BaseException] = [] - - def use_endpoint() -> None: - try: - endpoint.progress() - except BaseException as error: - errors.append(error) - - thread = threading.Thread(target=use_endpoint) - thread.start() - thread.join() - self.assertEqual(len(errors), 1) - self.assertIsInstance(errors[0], UnsupportedError) - endpoint.close() - - def test_caller_serialized_mode_allows_sequential_thread_migration(self) -> None: - source = Endpoint("source", thread_mode="caller_serialized") - destination = Endpoint("destination", thread_mode=ThreadMode.CALLER_SERIALIZED) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - work = source.write(plan) - - self.assertIs(source._state_lock, _api._NOOP_LOCK) - self.assertIsNone(work._lock) - self.assertIsNone(work._owner_thread) - failures: list[BaseException] = [] - external_lock = threading.Lock() - - def observe_and_close() -> None: - try: - with external_lock: - self.assertEqual(work.state, WorkState.COMPLETED) - work.close() - source.progress() - except BaseException as error: - failures.append(error) - - thread = threading.Thread(target=observe_and_close) - thread.start() - thread.join(2) - self.assertFalse(thread.is_alive()) - self.assertEqual(failures, []) - self.assertTrue(work.closed) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_single_indexed_success_never_enters_core_noop_lock(self) -> None: - source = Endpoint("source", thread_mode=ThreadMode.SINGLE) - destination = Endpoint("destination") - local = source.register(bytearray(b"abcdefgh"), name="buffer") - remote = destination.register(bytearray(8), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 4, stride=4, count=2)], - remote=[peer.region("buffer", 0, 4, stride=4, count=2)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[1], - remote_indices=[0], - ) - - with mock.patch.object( - _api._NoopLock, - "__enter__", - side_effect=AssertionError("successful indexed path entered _NoopLock"), - ): - work = selection.submit() - self.assertEqual(work.state, WorkState.COMPLETED) - work.close() - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_caller_serialized_exact_slots_migrate_without_core_locks(self) -> None: - source = Endpoint( - "source", - backend="test-request-slot-notification", - thread_mode="caller_serialized", - ) - destination = Endpoint( - "destination", - backend="test-request-slot-notification", - thread_mode="caller_serialized", - ) - local = source.register(bytearray(b"data"), name="buffer") - remote = destination.register(bytearray(4), name="buffer") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("buffer")]) - bound = plan.bind(TransferOp.WRITE) - fixed = bound.create_slot() - notification = bound.create_slot(request_local_notifications=True) - - self.assertIsNone(fixed._lock) - self.assertIsNone(fixed._owner_thread) - self.assertIs(fixed._fast_execute, fixed._backend_start_and_poll) - self.assertIsNone(notification._lock) - self.assertIsNone(notification._owner_thread) - self.assertIs( - notification._notification_fast_execute, - notification._backend_start_and_poll_with_notification, - ) - - external_lock = threading.Lock() - - def on_new_thread(function): - results = [] - failures: list[BaseException] = [] - - def invoke() -> None: - try: - with external_lock: - results.append(function()) - except BaseException as error: - failures.append(error) - - thread = threading.Thread(target=invoke) - thread.start() - thread.join(2) - self.assertFalse(thread.is_alive()) - self.assertEqual(failures, []) - self.assertEqual(len(results), 1) - return results[0] - - with mock.patch.object( - _api, - "_GET_IDENT", - side_effect=AssertionError( - "caller-serialized slot queried an unowned thread identity" - ), - ): - self.assertEqual(on_new_thread(fixed.execute), 1) - fixed_backend = source._backend.request_slots[0] - fixed_backend.next_state = WorkState.RUNNING - self.assertFalse(on_new_thread(fixed.start_and_test)) - fixed_backend.poll_states = [WorkState.COMPLETED] - self.assertTrue(on_new_thread(fixed.test_and_recycle)) - - self.assertEqual( - on_new_thread(lambda: notification.execute(notification=b"first")), - 1, - ) - notification_backend = source._backend.request_slots[1] - notification_backend.next_state = WorkState.RUNNING - self.assertFalse( - on_new_thread(lambda: notification.start_and_test(notification=None)) - ) - notification_backend.poll_states = [WorkState.COMPLETED] - self.assertTrue(on_new_thread(notification.test_and_recycle)) - - self.assertIs(fixed._fast_execute, fixed._backend_start_and_poll) - self.assertIs( - notification._notification_fast_execute, - notification._backend_start_and_poll_with_notification, - ) - notification.close() - fixed.close() - bound.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_metadata_rejects_unknown_fields(self) -> None: - endpoint = Endpoint("endpoint") - encoded = bytes(endpoint.export_metadata()) - corrupted = encoded[:-1] + b',"unknown":1}' - with self.assertRaises(InvalidMetadataError): - EndpointMetadata.from_bytes(corrupted) - overlong_integer = encoded.replace(b'"major":1', b'"major":' + b"9" * 4301, 1) - with self.assertRaises(InvalidMetadataError): - EndpointMetadata.from_bytes(overlong_integer) - with self.assertRaises(InvalidMetadataError): - EndpointMetadata.from_bytes(b"[" * 2000 + b"0" + b"]" * 2000) - endpoint.close() - - def test_metadata_wire_minor_extensions_and_bounds(self) -> None: - endpoint = Endpoint("endpoint") - current = endpoint.export_metadata() - future_optional = EndpointMetadata( - wire_major=current.wire_major, - wire_minor=current.wire_minor + 1, - required_features=(), - backend=current.backend, - endpoint_id=current.endpoint_id, - name=current.name, - incarnation=current.incarnation, - backend_payload=current.backend_payload, - registrations=current.registrations, - extensions={"future.optional": b"opaque"}, - ) - decoded = EndpointMetadata.from_bytes(bytes(future_optional)) - self.assertEqual(decoded.wire_minor, current.wire_minor + 1) - self.assertEqual(decoded.extensions, (("future.optional", b"opaque"),)) - - with self.assertRaisesRegex(InvalidMetadataError, "required"): - EndpointMetadata( - wire_major=current.wire_major, - wire_minor=current.wire_minor + 1, - required_features=("future.required",), - backend=current.backend, - endpoint_id=current.endpoint_id, - name=current.name, - incarnation=current.incarnation, - backend_payload=current.backend_payload, - registrations=current.registrations, - ) - with self.assertRaisesRegex(InvalidMetadataError, "size limit"): - EndpointMetadata.from_bytes(b" " * (EndpointMetadata.MAX_ENCODED_BYTES + 1)) - - multidimensional = memoryview(bytearray(b" ")).cast("B", (1, 4)) - self.assertEqual(len(multidimensional), 1) - self.assertEqual(multidimensional.nbytes, 4) - with ( - mock.patch.object(EndpointMetadata, "MAX_ENCODED_BYTES", 3), - self.assertRaisesRegex(InvalidMetadataError, "size limit"), - ): - EndpointMetadata.from_bytes(multidimensional) - - registration = RegistrationMetadata( - registration_id="r", - name="n", - nbytes=1, - device="cuda", - memory_type="m", - backend_descriptor=b"", - ) - with ( - mock.patch.object(EndpointMetadata, "MAX_STRING_BYTES", 3), - self.assertRaisesRegex(InvalidMetadataError, "registration device"), - ): - EndpointMetadata( - wire_major=current.wire_major, - wire_minor=current.wire_minor, - required_features=(), - backend="ref", - endpoint_id="id", - name="n", - incarnation="i", - backend_payload=b"", - registrations=(registration,), - ) - - with ( - mock.patch.object(EndpointMetadata, "MAX_EXTENSION_BYTES", 1), - self.assertRaisesRegex(InvalidMetadataError, "extension.*size limit"), - ): - EndpointMetadata( - wire_major=current.wire_major, - wire_minor=current.wire_minor, - required_features=(), - backend=current.backend, - endpoint_id=current.endpoint_id, - name=current.name, - incarnation=current.incarnation, - backend_payload=current.backend_payload, - registrations=current.registrations, - extensions={"optional": b"xx"}, - ) - - with self.assertRaisesRegex(InvalidMetadataError, "valid Unicode"): - EndpointMetadata( - wire_major=current.wire_major, - wire_minor=current.wire_minor, - required_features=(), - backend=current.backend, - endpoint_id=current.endpoint_id, - name="invalid-\ud800", - incarnation=current.incarnation, - backend_payload=current.backend_payload, - registrations=current.registrations, - ) - endpoint.close() - - def test_core_has_no_exception_unsafe_manual_lock_acquisition(self) -> None: - tree = ast.parse(Path(_api.__file__).read_text(encoding="utf-8")) - acquisitions = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "acquire" - ] - self.assertEqual(acquisitions, []) - - def test_first_publication_interruptions_roll_back_every_resource_family( - self, - ) -> None: - class InterruptAfterAdd(set): - def __init__(self, values=()) -> None: - super().__init__(values) - self.armed = True - self.last = None - - def add(self, value) -> None: - super().add(value) - self.last = value - if self.armed: - self.armed = False - raise KeyboardInterrupt("post-add publication cut") - - class InterruptAfterSet(dict): - def __init__(self, values=()) -> None: - super().__init__(values) - self.armed = True - self.last = None - - def __setitem__(self, key, value) -> None: - super().__setitem__(key, value) - self.last = value - if self.armed: - self.armed = False - raise KeyboardInterrupt("post-set publication cut") - - endpoint = Endpoint("registration-publication") - for attribute, name in ( - ("_registrations", "primary"), - ("_registration_names", "name-index"), - ): - collection = InterruptAfterSet(getattr(endpoint, attribute)) - setattr(endpoint, attribute, collection) - with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): - endpoint.register(bytearray(4), name=name) - self.assertTrue(collection.last.closed) - self.assertEqual(endpoint.open_registrations, ()) - self.assertFalse(endpoint._registration_names) - endpoint.close() - - source = Endpoint("source-publication") - destination = Endpoint("destination-publication") - local = source.register(bytearray(b"abcd"), name="local") - remote = destination.register(bytearray(4), name="remote") - metadata = bytes(destination.export_metadata()) - - peer_publication = InterruptAfterSet(source._peers) - source._peers = peer_publication - with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): - source.import_peer(metadata) - self.assertTrue(peer_publication.last.closed) - self.assertEqual(source.open_peers, ()) - peer = source.import_peer(metadata) - - plan_publication = InterruptAfterAdd(source._plans) - source._plans = plan_publication - with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): - source.prepare(local=[local.region()], remote=[peer.region("remote")]) - self.assertTrue(plan_publication.last.closed) - self.assertEqual(source.open_plans, ()) - - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("remote", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - ) - for public_first in (True, False): - if public_first: - publication = InterruptAfterAdd(source._bound_transfers) - source._bound_transfers = publication - else: - publication = InterruptAfterAdd(plan._bindings) - plan._bindings = publication - with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): - selection.bind() - self.assertTrue(publication.last.closed) - self.assertEqual(source.open_bound_transfers, ()) - self.assertFalse(plan._bindings) - - work_publication = InterruptAfterAdd(source._works) - source._works = work_publication - token = object() - with self.assertRaisesRegex(KeyboardInterrupt, "publication cut"): - source.send_notification(peer, b"not-sent", recovery_token=token) - self.assertTrue(work_publication.last.closed) - self.assertFalse(work_publication.last._submitting) - self.assertIsNone(work_publication.last.recovery_token) - self.assertEqual(source.open_works, ()) - - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_published_work_is_observation_gated_until_submit_reconciles(self) -> None: - for indexed in (False, True): - with self.subTest(indexed=indexed): - source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) - destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) - local = source.register(bytearray(b"abcd"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("remote", 0, 1, stride=1, count=4)], - indexed=indexed, - ) - backend_work = _CountingPendingWork() - entered = threading.Event() - release = threading.Event() - - def blocked_submit(*args, **kwargs) -> None: - kwargs["adopt_work"](backend_work) - entered.set() - release.wait(2) - - if indexed: - source._backend.submit_prevalidated = blocked_submit - selection = plan.select_indices( - TransferOp.WRITE, - local_indices=[0], - remote_indices=[0], - ) - submit = selection.submit - else: - source._backend.submit = blocked_submit - submit = lambda *, recovery_token: source.write( - plan, recovery_token=recovery_token - ) - token = object() - returned: list[Work] = [] - submit_thread = threading.Thread( - target=lambda: returned.append(submit(recovery_token=token)) - ) - submit_thread.start() - self.assertTrue(entered.wait(1)) - try: - (work,) = source.open_works - self.assertIs(work.recovery_token, token) - observations = ( - lambda: work.state, - lambda: work.error, - work.test, - lambda: work.poll_bounded(max_polls=1, timeout_ns=None), - work.cancel, - lambda: work.completion_event, - work.close, - ) - for observation in observations: - with self.assertRaisesRegex( - BusyError, "submission is in progress" - ): - observation() - self.assertEqual(backend_work.state_queries, 0) - with self.assertRaisesRegex(BusyError, "submissions in progress"): - plan.close() - finally: - release.set() - submit_thread.join(2) - self.assertFalse(submit_thread.is_alive()) - self.assertEqual(returned, [work]) - self.assertFalse(work._submitting) - backend_work._state = WorkState.COMPLETED - self.assertEqual(work.state, WorkState.COMPLETED) - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_indexed_publication_is_owned_before_lock_exit_and_every_graph_edge( - self, - ) -> None: - class InterruptAfterAdd(set): - def __init__(self, values=()) -> None: - super().__init__(values) - self.armed = True - self.last = None - - def add(self, value) -> None: - super().add(value) - self.last = value - if self.armed: - self.armed = False - raise KeyboardInterrupt("indexed graph publication cut") - - class InterruptAfterExit: - def __init__(self, lock) -> None: - self.lock = lock - self.armed = True - - def __enter__(self): - self.lock.__enter__() - return self - - def __exit__(self, exc_type, exc_value, traceback): - result = self.lock.__exit__(exc_type, exc_value, traceback) - if self.armed: - self.armed = False - raise KeyboardInterrupt("indexed state-lock exit cut") - return result - - source = Endpoint("source", thread_mode=ThreadMode.MULTIPLE) - destination = Endpoint("destination", thread_mode=ThreadMode.MULTIPLE) - local = source.register(bytearray(b"abcd"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare( - local=[local.region(0, 1, stride=1, count=4)], - remote=[peer.region("remote", 0, 1, stride=1, count=4)], - indexed=True, - ) - selection = plan.select_indices( - TransferOp.WRITE, local_indices=[0], remote_indices=[0] - ) - backend_entries = 0 - - def unexpected_backend_entry(*args, **kwargs) -> None: - nonlocal backend_entries - backend_entries += 1 - raise AssertionError("publication interruption reached provider") - - source._backend.submit_prevalidated = unexpected_backend_entry - original_state_lock = source._state_lock - source._state_lock = InterruptAfterExit(original_state_lock) - try: - with self.assertRaisesRegex(KeyboardInterrupt, "state-lock exit cut"): - selection.submit(recovery_token=object()) - finally: - source._state_lock = original_state_lock - self.assertEqual(source.open_works, ()) - self.assertFalse(plan._submitting) - self.assertFalse(plan._works) - - for owner, attribute in ( - (source, "_works"), - (plan, "_submitting"), - (plan, "_works"), - ): - with self.subTest(edge=f"{type(owner).__name__}.{attribute}"): - original = getattr(owner, attribute) - publication = InterruptAfterAdd(original) - setattr(owner, attribute, publication) - try: - with self.assertRaisesRegex( - KeyboardInterrupt, "graph publication cut" - ): - selection.submit(recovery_token=object()) - self.assertIsNotNone(publication.last) - self.assertTrue(publication.last.closed) - self.assertFalse(publication.last._submitting) - self.assertIsNone(publication.last.recovery_token) - finally: - setattr(owner, attribute, set(publication)) - self.assertEqual(source.open_works, ()) - self.assertFalse(plan._submitting) - self.assertFalse(plan._works) - - self.assertEqual(backend_entries, 0) - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_endpoint_post_adoption_initialization_cleanup_is_complete(self) -> None: - for cut in ("capabilities", "dispatch", "endpoint-info"): - with self.subTest(cut=cut): - instances = [] - - class InitializationBackend(_ReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.close_calls = 0 - - @property - def capabilities(self): - if cut == "capabilities": - raise KeyboardInterrupt("capabilities cut") - return super().capabilities - - def __getattribute__(self, name): - if name == "send_notification_sync": - current_cut = object.__getattribute__(self, "_test_cut") - if current_cut == "dispatch": - raise KeyboardInterrupt("dispatch cut") - return super().__getattribute__(name) - - def close(self) -> None: - self.close_calls += 1 - return super().close() - - def factory(*args, **kwargs): - implementation = InitializationBackend(*args, **kwargs) - implementation._test_cut = cut - instances.append(implementation) - return implementation - - backend_name = f"test-endpoint-init-cleanup-{cut}" - register_backend(backend_name, factory, replace=True) - if cut == "endpoint-info": - with ( - mock.patch.object( - _api, - "EndpointInfo", - side_effect=KeyboardInterrupt("EndpointInfo cut"), - ), - self.assertRaisesRegex(KeyboardInterrupt, "EndpointInfo cut"), - ): - Endpoint("interrupted", backend=backend_name) - else: - with self.assertRaisesRegex(KeyboardInterrupt, f"{cut} cut"): - Endpoint("interrupted", backend=backend_name) - self.assertEqual(len(instances), 1) - self.assertEqual(instances[0].close_calls, 1) - self.assertTrue(instances[0]._closed) - - def test_endpoint_failed_initialization_exposes_retryable_cleanup_owner( - self, - ) -> None: - marker = object() - instances = [] - - class HeldInitializationBackend(_ReferenceBackend): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.close_calls = 0 - - def close(self): - self.close_calls += 1 - super().close() - return marker if self.close_calls == 1 else None - - def factory(*args, **kwargs): - implementation = HeldInitializationBackend(*args, **kwargs) - instances.append(implementation) - return implementation - - backend_name = "test-endpoint-init-held-cleanup" - register_backend(backend_name, factory, replace=True) - with ( - mock.patch.object( - _api, - "EndpointInfo", - side_effect=KeyboardInterrupt("EndpointInfo cut"), - ), - self.assertRaisesRegex(BackendFailureError, "recovery_endpoint") as context, - ): - Endpoint("interrupted", backend=backend_name) - - recovery_endpoint = context.exception.recovery_endpoint - self.assertIs(recovery_endpoint._backend, instances[0]) - self.assertFalse(recovery_endpoint.closed) - self.assertTrue(recovery_endpoint._creation_failed) - self.assertEqual(instances[0].close_calls, 1) - with self.assertRaisesRegex(BackendFailureError, "creation cleanup only"): - recovery_endpoint.register(bytearray(1)) - recovery_endpoint.close() - self.assertTrue(recovery_endpoint.closed) - self.assertFalse(recovery_endpoint._creation_failed) - self.assertEqual(instances[0].close_calls, 2) - - def test_submit_success_epilogues_keep_noop_lock_fast_branches(self) -> None: - source = Path(_api.__file__).read_text(encoding="utf-8") - generic = source[source.index(" def _submit(\n") :] - generic = generic[: generic.index(" def _validate_regions(\n")] - bound = source[source.index(" def _submit_bound(\n") :] - bound = bound[: bound.index(" def send_notification(\n")] - self.assertIn("if state_lock is _NOOP_LOCK:", generic) - self.assertIn("if state_lock is _NOOP_LOCK:", bound) - poll = source[source.index(" def poll_notification_batches(\n") :] - poll = poll[: poll.index(" def progress(")] - conditional = poll.index("if max_items is not None:") - payload_count = poll.index( - "payload_count = sum(len(batch.payloads) for batch in batches)" - ) - self.assertLess(conditional, payload_count) - - def test_teardown_requires_exact_none_before_core_commit(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - work = source.write(plan) - self.assertEqual(work.state, WorkState.COMPLETED) - - marker = object() - - def expect_retryable(resource, owner, attribute, original, *args) -> None: - def close_then_return_marker(*call_args, **call_kwargs): - original(*call_args, **call_kwargs) - return marker - - setattr(owner, attribute, close_then_return_marker) - with self.assertRaisesRegex(BackendFailureError, "must return None"): - resource.close() - self.assertFalse(resource.closed) - setattr(owner, attribute, original) - resource.close() - - backend_work = work._backend_work - expect_retryable(work, backend_work, "close", backend_work.close) - expect_retryable( - plan, - source._backend, - "release_plan", - source._backend.release_plan, - plan._handle, - ) - expect_retryable( - peer, - source._backend, - "release_peer", - source._backend.release_peer, - peer._handle, - ) - expect_retryable( - local, - source._backend, - "deregister", - source._backend.deregister, - local._handle, - ) - remote.close() - original_close = source._backend.close - expect_retryable(source, source._backend, "close", original_close) - destination.close() - - def test_work_completion_event_is_capability_gated(self) -> None: - source = Endpoint("source") - destination = Endpoint("destination") - local = source.register(bytearray(b"data"), name="local") - remote = destination.register(bytearray(4), name="remote") - peer = source.import_peer(bytes(destination.export_metadata())) - plan = source.prepare(local=[local.region()], remote=[peer.region("remote")]) - work = source.write(plan) - with mock.patch.object( - type(work._backend_work), - "completion_event", - new_callable=mock.PropertyMock, - side_effect=AssertionError("disabled capability entered provider"), - ): - self.assertIsNone(work.completion_event) - work.close() - plan.close() - peer.close() - local.close() - remote.close() - source.close() - destination.close() - - def test_notification_batch_exact_type_fast_path_skips_payload_rescan(self) -> None: - endpoint = Endpoint("receiver") - payloads = (b"one", b"two") - batch = BackendNotificationBatch(payloads) - endpoint._backend.poll_notification_batches = lambda max_items: [batch] - with mock.patch.object( - _types, - "_immutable_notification_payloads", - side_effect=AssertionError("validated backend tuple was rescanned"), - ): - (public,) = endpoint.poll_notification_batches() - self.assertIs(public.payloads, payloads) - - class BackendBatchSubclass(BackendNotificationBatch): - pass - - endpoint._backend.poll_notification_batches = lambda max_items: [ - BackendBatchSubclass((b"subclass",)) - ] - with self.assertRaisesRegex(BackendFailureError, "must be a list"): - endpoint.poll_notification_batches() - endpoint.close() - - def test_provider_objects_redact_owners_and_capability_descriptors(self) -> None: - secret = b"provider-capability-secret" - owner = {"secret": secret} - registered = RegisteredMemory("id", "name", owner, 1, 1, "cpu", "host") - imported = ImportedRegistration( - owner, - "id", - "name", - 1, - "cpu", - "host", - secret, - ) - raw = ImportedRawMemory(owner, 1, 1, "cpu", "host") - region = BackendRegion( - owner=owner, - offset=0, - nbytes=1, - device="cpu", - memory_type="host", - backend_descriptor=secret, - ) - - for value in (registered, imported, raw, region): - self.assertNotIn("provider-capability-secret", repr(value)) + self.a.submit("WRITE", self.local, [2], self.remote, [0]) + with self.assertRaises(ValueError): + self.a.submit("WRITE", self.local, [], self.remote, []) + with self.assertRaises(ValueError): + self.a.register_tensor(torch.zeros(2, 2).t()) if __name__ == "__main__": - run_tests() + unittest.main() diff --git a/test/dynamo/test_trace_rules.py b/test/dynamo/test_trace_rules.py index 1d4d7206e1113..d7c27c6efd291 100644 --- a/test/dynamo/test_trace_rules.py +++ b/test/dynamo/test_trace_rules.py @@ -191,7 +191,6 @@ def _is_allowed_module_prefix(obj): "torch._C._distributed_autograd", "torch._C._distributed_c10d", "torch._C._distributed_rpc", - "torch._C._distributed_transfer", "torch._C._functorch", "torch._C._monitor", "torch._C._nvtx", diff --git a/torch/_C/_distributed_transfer.pyi b/torch/_C/_distributed_transfer.pyi deleted file mode 100644 index b279648311126..0000000000000 --- a/torch/_C/_distributed_transfer.pyi +++ /dev/null @@ -1,7 +0,0 @@ -# This module is defined in torch/csrc/distributed/transfer/python_bindings.cpp - -def _snapshot_and_analyze_int32_indices( - local_indices: object, - remote_indices: object, - destination_is_remote: bool, -) -> tuple[bytes, int, int, int, int, int, int, bool]: ... diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp index bd384cacc35e6..200dcebdf572d 100644 --- a/torch/csrc/Module.cpp +++ b/torch/csrc/Module.cpp @@ -73,7 +73,6 @@ #include #include #include -#include #include #include #include @@ -2670,7 +2669,6 @@ PyObject* initModule() { #endif torch::distributed::initPlacementBindings(module); - torch::distributed::transfer::initPythonBindings(module); torch::initTensorIteratorBindings(module); diff --git a/torch/csrc/distributed/transfer/python_bindings.cpp b/torch/csrc/distributed/transfer/python_bindings.cpp deleted file mode 100644 index 4969acd80cf71..0000000000000 --- a/torch/csrc/distributed/transfer/python_bindings.cpp +++ /dev/null @@ -1,432 +0,0 @@ -#include - -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace py = pybind11; - -namespace torch::distributed::transfer { -namespace { - -constexpr std::size_t kTinyDuplicateThreshold = 32; -constexpr std::uint64_t kMaximumBitmapValueCount = 1ULL << 24; -constexpr std::uint64_t kMaximumBitmapValuesPerIndex = 16; - -class BufferView final { - public: - BufferView(PyObject* object, const char* name) : name_(name) { - if (PyObject_GetBuffer(object, &view_, PyBUF_FORMAT | PyBUF_STRIDES) != 0) { - throw py::error_already_set(); - } - acquired_ = true; - try { - validate(); - } catch (...) { - release(); - throw; - } - } - - BufferView(const BufferView&) = delete; - BufferView& operator=(const BufferView&) = delete; - - ~BufferView() { - release(); - } - - const void* data() const { - return view_.buf; - } - - Py_ssize_t nbytes() const { - return view_.len; - } - - std::size_t count() const { - return static_cast(view_.len) / sizeof(std::int32_t); - } - - private: - static char nativeEndianPrefix() { - const std::uint16_t value = 1; - std::uint8_t first_byte = 0; - std::memcpy(&first_byte, &value, sizeof(first_byte)); - return first_byte == 1 ? '<' : '>'; - } - - static bool isNativeInt32Format(const char* format) { - if (format == nullptr) { - return false; - } - if (std::strcmp(format, "i") == 0 || std::strcmp(format, "@i") == 0 || - std::strcmp(format, "=i") == 0) { - return true; - } - return format[0] == nativeEndianPrefix() && - std::strcmp(format + 1, "i") == 0; - } - - void validate() const { - if (view_.ndim != 1 || view_.itemsize != sizeof(std::int32_t) || - view_.len < 0 || view_.len % sizeof(std::int32_t) != 0 || - !isNativeInt32Format(view_.format)) { - throw py::value_error( - std::string(name_) + - " must be a one-dimensional C-contiguous native-int32 buffer"); - } - const auto count = view_.len / sizeof(std::int32_t); - if (view_.buf == nullptr || view_.shape == nullptr || - view_.strides == nullptr || view_.shape[0] != count || - (view_.suboffsets != nullptr && view_.suboffsets[0] >= 0)) { - throw py::value_error( - std::string(name_) + " has an inconsistent PEP 3118 shape"); - } - if (PyBuffer_IsContiguous(&view_, 'C') == 0) { - throw py::value_error( - std::string(name_) + - " must be a one-dimensional C-contiguous native-int32 buffer"); - } - if (count == 0) { - throw py::value_error(std::string(name_) + " must be non-empty"); - } - if (static_cast(view_.len) > - std::numeric_limits::max()) { - throw std::overflow_error( - std::string(name_) + " is too large for this platform"); - } - } - - void release() { - if (acquired_) { - PyBuffer_Release(&view_); - acquired_ = false; - } - } - - const char* name_; - Py_buffer view_{}; - bool acquired_{false}; -}; - -struct IndexAnalysis final { - std::int32_t local_minimum; - std::int32_t local_maximum; - std::int32_t remote_minimum; - std::int32_t remote_maximum; - bool destination_has_duplicate; -}; - -struct BoundsAnalysis final { - std::int32_t minimum; - std::int32_t maximum; -}; - -struct DestinationAnalysis final { - BoundsAnalysis bounds; - bool strictly_monotonic; -}; - -std::int32_t loadInt32(const char* data, std::size_t index) { - std::int32_t value = 0; - std::memcpy( - &value, data + index * sizeof(std::int32_t), sizeof(std::int32_t)); - return value; -} - -std::uint64_t mixKey(std::uint64_t value) { - value += 0x9e3779b97f4a7c15ULL; - value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; - value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; - return value ^ (value >> 31); -} - -bool hasTinyDuplicate(const char* data, std::size_t count) { - for (std::size_t index = 1; index < count; ++index) { - const auto value = loadInt32(data, index); - for (std::size_t prior = 0; prior < index; ++prior) { - if (loadInt32(data, prior) == value) { - return true; - } - } - } - return false; -} - -bool hasBitmapDuplicate( - const char* data, - std::size_t count, - std::int32_t minimum, - std::uint64_t value_count) { - std::vector bitmap( - static_cast((value_count + 63) / 64), 0); - for (std::size_t index = 0; index < count; ++index) { - const auto offset = static_cast( - static_cast(loadInt32(data, index)) - minimum); - const auto word = static_cast(offset / 64); - const auto mask = std::uint64_t{1} << (offset % 64); - if ((bitmap[word] & mask) != 0) { - return true; - } - bitmap[word] |= mask; - } - return false; -} - -bool hasFlatHashDuplicate(const char* data, std::size_t count) { - if (count > std::numeric_limits::max() / 2) { - throw std::overflow_error("int32 index count exceeds hash-table capacity"); - } - const auto requested_capacity = count * 2; - std::size_t capacity = 1; - while (capacity < requested_capacity) { - if (capacity > std::numeric_limits::max() / 2) { - throw std::overflow_error("int32 index hash-table capacity overflow"); - } - capacity *= 2; - } - std::vector table(capacity, 0); - const auto mask = capacity - 1; - for (std::size_t index = 0; index < count; ++index) { - const auto key = static_cast( - static_cast(loadInt32(data, index))) + - 1; - auto slot = static_cast(mixKey(key)) & mask; - while (table[slot] != 0) { - if (table[slot] == key) { - return true; - } - slot = (slot + 1) & mask; - } - table[slot] = key; - } - return false; -} - -bool hasDuplicate( - const char* data, - std::size_t count, - std::int32_t minimum, - std::int32_t maximum, - bool strictly_monotonic) { - if (count < 2 || strictly_monotonic) { - return false; - } - const auto value_count = - static_cast(static_cast(maximum) - minimum) + - 1; - if (value_count < count) { - return true; - } - if (count <= kTinyDuplicateThreshold) { - return hasTinyDuplicate(data, count); - } - const auto dense_enough = - count >= (value_count + kMaximumBitmapValuesPerIndex - 1) / - kMaximumBitmapValuesPerIndex; - if (value_count <= kMaximumBitmapValueCount && dense_enough) { - return hasBitmapDuplicate(data, count, minimum, value_count); - } - return hasFlatHashDuplicate(data, count); -} - -BoundsAnalysis analyzeBounds(const char* data, std::size_t count) { - const auto first = loadInt32(data, 0); - BoundsAnalysis result{first, first}; - for (std::size_t index = 1; index < count; ++index) { - const auto value = loadInt32(data, index); - result.minimum = std::min(result.minimum, value); - result.maximum = std::max(result.maximum, value); - } - return result; -} - -DestinationAnalysis finishNonmonotonicDestination( - const char* data, - std::size_t count, - std::size_t suffix_index, - BoundsAnalysis bounds) { - if (suffix_index < count) { - const auto suffix = analyzeBounds( - data + suffix_index * sizeof(std::int32_t), count - suffix_index); - bounds.minimum = std::min(bounds.minimum, suffix.minimum); - bounds.maximum = std::max(bounds.maximum, suffix.maximum); - } - return {bounds, false}; -} - -DestinationAnalysis analyzeDestination(const char* data, std::size_t count) { - // Real framework selections are commonly ordered. Prove their direction - // with a short scalar probe, then let the compiler vectorize the long tail's - // order and bounds reductions. Non-monotonic inputs leave the probe quickly - // and retain the same exact duplicate-analysis fallbacks. - constexpr std::size_t kMonotonicProbeElementCount = 8; - const auto first = loadInt32(data, 0); - DestinationAnalysis result{{first, first}, true}; - if (count == 1) { - return result; - } - - auto previous = loadInt32(data, 1); - if (previous == first) { - return finishNonmonotonicDestination(data, count, 2, result.bounds); - } - - const auto probe_end = std::min(count, kMonotonicProbeElementCount); - std::size_t index = 2; - if (previous > first) { - result.bounds.maximum = previous; - for (; index < probe_end; ++index) { - const auto value = loadInt32(data, index); - if (value <= previous) { - result.bounds.minimum = std::min(result.bounds.minimum, value); - return finishNonmonotonicDestination( - data, count, index + 1, result.bounds); - } - result.bounds.maximum = value; - previous = value; - } - - auto minimum = result.bounds.minimum; - auto maximum = result.bounds.maximum; - bool strictly_monotonic = true; - for (; index < count; ++index) { - const auto value = loadInt32(data, index); - minimum = std::min(minimum, value); - maximum = std::max(maximum, value); - strictly_monotonic &= loadInt32(data, index - 1) < value; - } - return {{minimum, maximum}, strictly_monotonic}; - } - - result.bounds.minimum = previous; - for (; index < probe_end; ++index) { - const auto value = loadInt32(data, index); - if (value >= previous) { - result.bounds.maximum = std::max(result.bounds.maximum, value); - return finishNonmonotonicDestination( - data, count, index + 1, result.bounds); - } - result.bounds.minimum = value; - previous = value; - } - - auto minimum = result.bounds.minimum; - auto maximum = result.bounds.maximum; - bool strictly_monotonic = true; - for (; index < count; ++index) { - const auto value = loadInt32(data, index); - minimum = std::min(minimum, value); - maximum = std::max(maximum, value); - strictly_monotonic &= loadInt32(data, index - 1) > value; - } - return {{minimum, maximum}, strictly_monotonic}; -} - -IndexAnalysis analyzeSnapshot( - const char* data, - std::size_t local_count, - std::size_t remote_count, - bool destination_is_remote) { - const auto* remote = data + local_count * sizeof(std::int32_t); - BoundsAnalysis local_bounds{}; - BoundsAnalysis remote_bounds{}; - DestinationAnalysis destination{}; - if (destination_is_remote) { - local_bounds = analyzeBounds(data, local_count); - destination = analyzeDestination(remote, remote_count); - remote_bounds = destination.bounds; - } else { - destination = analyzeDestination(data, local_count); - local_bounds = destination.bounds; - remote_bounds = analyzeBounds(remote, remote_count); - } - IndexAnalysis result{ - local_bounds.minimum, - local_bounds.maximum, - remote_bounds.minimum, - remote_bounds.maximum, - false}; - const auto* destination_data = destination_is_remote ? remote : data; - const auto destination_count = - destination_is_remote ? remote_count : local_count; - result.destination_has_duplicate = hasDuplicate( - destination_data, - destination_count, - destination.bounds.minimum, - destination.bounds.maximum, - destination.strictly_monotonic); - return result; -} - -py::tuple snapshotAndAnalyzeInt32Indices( - py::object local_indices, - py::object remote_indices, - bool destination_is_remote) { - py::bytes snapshot; - std::size_t local_count = 0; - std::size_t remote_count = 0; - { - BufferView local(local_indices.ptr(), "local_indices"); - BufferView remote(remote_indices.ptr(), "remote_indices"); - if (local.nbytes() > PY_SSIZE_T_MAX - remote.nbytes()) { - throw std::overflow_error("combined int32 index snapshot is too large"); - } - const auto snapshot_nbytes = local.nbytes() + remote.nbytes(); - PyObject* raw_snapshot = - PyBytes_FromStringAndSize(nullptr, snapshot_nbytes); - if (raw_snapshot == nullptr) { - throw py::error_already_set(); - } - snapshot = py::reinterpret_steal(raw_snapshot); - auto* snapshot_data = PyBytes_AS_STRING(snapshot.ptr()); - std::memcpy(snapshot_data, local.data(), local.nbytes()); - std::memcpy(snapshot_data + local.nbytes(), remote.data(), remote.nbytes()); - local_count = local.count(); - remote_count = remote.count(); - } - - const auto* snapshot_data = PyBytes_AS_STRING(snapshot.ptr()); - IndexAnalysis analysis{}; - { - py::gil_scoped_release release; - analysis = analyzeSnapshot( - snapshot_data, local_count, remote_count, destination_is_remote); - } - return py::make_tuple( - std::move(snapshot), - local_count, - remote_count, - analysis.local_minimum, - analysis.local_maximum, - analysis.remote_minimum, - analysis.remote_maximum, - analysis.destination_has_duplicate); -} - -} // namespace - -void initPythonBindings(PyObject* module) { - auto root = py::handle(module).cast(); - auto transfer = root.def_submodule( - "_distributed_transfer", "Private distributed transfer helpers."); - transfer.def( - "_snapshot_and_analyze_int32_indices", - &snapshotAndAnalyzeInt32Indices, - py::arg("local_indices"), - py::arg("remote_indices"), - py::arg("destination_is_remote"), - R"doc( -Create one immutable packed snapshot of two native-int32 buffers and analyze it. - -This is a private implementation helper for ``torch.distributed._transfer``. -)doc"); -} - -} // namespace torch::distributed::transfer diff --git a/torch/csrc/distributed/transfer/python_bindings.h b/torch/csrc/distributed/transfer/python_bindings.h deleted file mode 100644 index 7518e726c9717..0000000000000 --- a/torch/csrc/distributed/transfer/python_bindings.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include - -namespace torch::distributed::transfer { - -void initPythonBindings(PyObject* module); - -} // namespace torch::distributed::transfer diff --git a/torch/distributed/_transfer/README.md b/torch/distributed/_transfer/README.md index 472e09c4dda0c..3b2392732c07d 100644 --- a/torch/distributed/_transfer/README.md +++ b/torch/distributed/_transfer/README.md @@ -1,814 +1,54 @@ -# Experimental endpoint transfer +# Transfer API — review MVP -`torch.distributed._transfer` is a private, backend-neutral data-plane API for -moving registered tensor or buffer bytes between independently managed -endpoints. It is intended for inference KV movement and live weight streaming, -where collective process-group semantics are the wrong abstraction. - -NIXL is the first external provider. PyTorch does not import, link, or otherwise -depend on NIXL. The in-process CPU `reference` backend exists for contract tests, -not performance. - -## Scope - -The layer owns: - -- allocation lifetime anchors and byte-range validation; -- versioned endpoint and named-registration metadata; -- peer, plan, work, binding, and persistent-request-slot resource ordering; -- READ, WRITE, indexed transfer, progress, cancellation, and notifications; -- typed errors and explicit backend capabilities; -- an experimental Python provider SPI. - -Frameworks continue to own discovery, routing, leases, cache/model versioning, -layout policy, scheduling, and failure recovery. This package is not a new -control plane and does not provide durable checkpointing. - -## Basic flow +Independent of ProcessGroup: move registered bytes between named peers using an +injected provider. No NIXL import/build dependency in Core. This private API is +experimental and deliberately not a production transport. ```python -from torch.distributed import _transfer - -source = _transfer.Endpoint("source", backend="reference") -destination = _transfer.Endpoint("destination", backend="reference") - -source_registration = source.register(source_tensor, name="kv") -destination_registration = destination.register(destination_tensor, name="kv") +import torch +from torch.distributed._transfer import Endpoint +from torch.distributed._transfer._reference import MemoryBackend -# The application transports these opaque bytes over its control plane. -destination_metadata = bytes(destination.export_metadata()) -destination_peer = source.import_peer(destination_metadata) - -plan = source.prepare( - local=[source_registration.region()], - remote=[destination_peer.region("kv")], -) -work = source.write(plan, notification=b"kv-visible") +source, target = torch.arange(8, dtype=torch.uint8), torch.zeros(8, dtype=torch.uint8) +a, b = Endpoint(MemoryBackend()), Endpoint(MemoryBackend()) +a.register_tensor(source) +b.register_tensor(target) +peer = a.connect(b.metadata()) # a real application exchanges this out of band +local = a.prepare([(source.data_ptr(), 4, 0, 4, 2)], kind="host") +remote = a.prepare([(target.data_ptr(), 4, 0, 4, 2)], kind="host", peer=peer) +work = a.submit("WRITE", local, [1, 0], remote, [0, 1]) work.wait() - -work.close() -plan.close() -destination_peer.close() -source_registration.close() -destination_registration.close() -source.close() -destination.close() -``` - -External NIXL registration is explicit and lazy: - -```python -from nixl.torch_transfer import register_torch_backend - -register_torch_backend() -endpoint = _transfer.Endpoint("worker", backend="nixl") -``` - -Importing `torch` or `nixl` alone does not activate the integration. - -## Provider construction ownership - -Existing providers remain source-compatible with the version-1 factory, which -returns a `Backend`. New native providers should register FactoryV2 explicitly: - -```python -def factory(endpoint_id, name, incarnation, progress_mode, thread_mode, options, - *, adopt_backend): - backend = ProviderBackend.new_unactivated(...) - adopt_backend(backend) # Core owns cleanup from this point - backend.activate_native() # may publish native identity or start threads - return None - -register_backend("provider", factory, factory_api_version=2) -``` - -A version-1 factory return is adopted only after Core verifies that it is a -`Backend`. A wrong-type return is foreign and unowned: Core raises `TypeError` -without invoking `close()` or any other method on that object. Version 1 cannot -close the interruption window between the factory return and this adoption; -new providers should use FactoryV2. - -Core advertises `BACKEND_FACTORY_API_VERSION` as the maximum backward-compatible -factory contract it understands. It never inspects a callable signature or -retries under another version. A version-2 provider must adopt exactly once -before activation and return `None`. Core retains an interruption-durable close -receipt and weak finalizer; a failed orphan is quarantined by exact -`(backend_name, endpoint_id)`, so an unresolved identity collision fails closed -without blocking unrelated endpoint IDs. - -Backend registration itself is retry-safe across a lost return. Re-registering -the same name with the identical callable object and exact factory API version -is an idempotent no-op under the registry lock. A different callable identity -or version still raises `BusyError`; providers cannot use retry recovery to -replace another owner's registration. - -`Endpoint.close()` does not call `progress()` implicitly. Applications using -manual progress own their drain policy; Core still rejects active Work and open -child resources before it starts retry-only backend teardown. - -## Memory model - -`Endpoint.register` accepts a contiguous dense Tensor, Storage, writable -C-contiguous buffer, or capability-gated `RawSpan`. It registers the exact -logical byte interval and never creates a contiguous staging copy. A Tensor -registration anchors its original Storage, so rebinding the Tensor with `set_` -does not invalidate the registered address. - -The caller must keep the underlying allocation address and extent unchanged -until the Registration is closed. In particular, Storage must not be resized -while a plan, Work, or TransferSlot can reference it. Core checks the anchored -address and size -before metadata export, plan creation, and submission, but that check is only a -diagnostic: Python cannot close the race with a concurrent resize after work is -submitted. Overlapping local registrations are rejected, including aliases of -the same Tensor, Storage, or writable buffer. - -Named `Peer.region` is the safe remote API. `Peer.unsafe_region` exists only for -legacy protocols that exchange addresses; the provider must prove that the raw -span is covered by imported remote metadata. Backends that cannot prove this -must reject it. - -A region can represent one interval or a compact strided layout: - -```python -region = registration.region( - offset=0, - nbytes=block_bytes, - stride=block_stride, - count=num_blocks, -) -``` - -The allocation extent is bounds checked; `nbytes * count` is the transfer size. - -## Indexed plans - -`prepare(..., indexed=True)` creates reusable local and remote descriptor -catalogs. Each Region contributes `count` logical blocks, or one block when -`count` is absent. Indices address the flattened region-major, block-minor -catalog, so a compact strided Region can expose many independently selectable -blocks without expanding the plan input. `TransferPlan.local_block_count` and -`remote_block_count` report the resulting bounds. - -Local and remote block counts and aggregate sizes may differ. Each submission -selects independent local and remote block-index sequences, whose selected byte -totals must match. Each selected block contributes its Region's `nbytes` value. -Only selected destination blocks participate in alias validation; duplicate or -physically overlapping destination selections are rejected. Omitting both index -sequences selects every logical block in each catalog. - -```python -plan = endpoint.prepare( - local=[local_registration.region(0, block_bytes, stride=64, count=8)], - remote=[remote_peer.region("blocks", 0, block_bytes, stride=64, count=32)], - indexed=True, -) -work = endpoint.write( - plan, - local_indices=[7, 2], - remote_indices=[19, 4], -) -``` - -Individual providers may impose a stricter mapping. The current NIXL provider -requires selected local and remote sequences to have equal lengths and matching -per-block sizes. - -### One-shot and reusable dynamic index selections - -Frameworks that produce a new packed block table per request should use the -one-shot plan API. Core snapshots and authenticates the indices, but does not -allocate a transient public `IndexSelection` or enter the endpoint state lock a -second time: - -```python -import numpy as np - -work, state = plan.submit_indices_and_poll( - _transfer.TransferOp.WRITE, - local_indices=np.asarray(local_block_ids, dtype=np.int32), - remote_indices=np.asarray(remote_block_ids, dtype=np.int32), - notification=request_tag, - max_polls=0, -) -``` - -`plan.submit_indices(...)` is the post-only form. Providers without fused -submission support use the same provider-neutral fallback; framework adapters -do not need a capability branch. - -When the same validated selection will be bound, reused, or turned into a -persistent slot, retain it explicitly: - -```python -selection = plan.select_indices( - _transfer.TransferOp.WRITE, - local_indices=np.asarray(local_block_ids, dtype=np.int32), - remote_indices=np.asarray(remote_block_ids, dtype=np.int32), -) -work, state = selection.submit_and_poll( - notification=request_tag, - max_polls=64, - timeout_ns=None, -) -while state not in { - _transfer.WorkState.COMPLETED, - _transfer.WorkState.FAILED, - _transfer.WorkState.CANCELLED, -}: - state = work.poll_bounded( - max_polls=4096, - timeout_ns=100_000, - timeout_check_interval=32, - ) -work.wait() # raises for FAILED/CANCELLED, otherwise returns immediately work.close() -``` - -Both `submit()` forms remain strictly post-only and are the correct QD>1 -primitive. A -provider may advertise `fused_prevalidated_submit_poll` and `fused_work_poll` -independently. The NIXL provider uses them to keep each bounded loop behind one -GIL release; providers without either capability use the source-compatible -composition. A zero initial poll budget posts only. On a non-fused provider its -returned Boolean is conservatively false because Core cannot recover a -post-derived native state without issuing a status observation. Boolean -`submit_indices_and_test`, `submit_and_test`, and `test_bounded` wrappers -remain available when a caller needs only terminal versus nonterminal. - -All bounded-poll counts, nanosecond budgets, and timeout-check intervals must -fit in a non-negative (or positive, where zero is not meaningful) signed -64-bit integer, matching the native provider ABI. A zero timeout still permits -exactly the first state observation even when the sampling interval is larger -than one; later clock checks occur at interval multiples. - -`IndexSelection` is immutable and belongs to the exact plan and operation used -to create it. Core copies both mutable, one-dimensional, C-contiguous native -`int32` inputs into one private packed buffer; NIXL reads that snapshot through -`numpy.frombuffer` without another copy. A general `Sequence` is still accepted: -Core applies `operator.index`, rejects booleans, validates bounds, and packs it -in one pass. Because this provider-neutral fast representation is signed -32-bit, an indexed catalog used through `select_indices` is limited to -2,147,483,648 logical entries. The format is deliberately native-endian so it -can be consumed directly by Torch and native providers; non-native-endian or -non-contiguous buffers take the safe sequence path. - -Uniform block size and full-catalog destination disjointness are plan facts. -Core can therefore validate byte totals without expanding regions and validate -destination uniqueness with one packed sort. Irregular or conservatively -ambiguous catalogs retain exact region expansion and physical-overlap checks. -The provider may omit its duplicate per-selection work only when Core's -authenticated token proves uniform geometry, full-catalog disjointness, and -unique destination indices *and* the provider independently proves the matching -full-catalog facts. - -Bind a selection when the same mapping is submitted repeatedly with a -request-local notification: - -```python -bound = selection.bind() -for request_tag in generation_tags: - work = bound.submit(notification=request_tag) - work.wait() - work.close() -bound.close() -``` - -The current NIXL raw binding updates or explicitly clears the notification at -every post. The Core reuse token identifies only the exact plan, operation, and -index snapshot; it never encodes provider-specific notification behavior. -Older NIXL bindings remain safe because the provider retires a cached request -before crossing between notification-present and notification-absent states. -Changing one non-empty tag to another still reuses the same request. Closing the -plan invalidates an otherwise live selection; each open Work or -selection-backed binding retains its packed snapshot until native access is -finished. - -## Bound submissions - -Latency-sensitive paths can bind one fully specified submission shape once and -reuse it without repeating plan ownership, storage stability, index, overlap, -notification, or CUDA-token validation: - -```python -bound = plan.bind( - _transfer.TransferOp.WRITE, - local_indices=[7, 2], - remote_indices=[19, 4], - cuda_ordering=_transfer.CudaOrderingMode.CALLER_READY, -) - -for _ in range(iterations): - work = bound.submit() - work.wait() - work.close() - -bound.close() -``` - -Every `submit()` creates a normal, independently owned `Work`; concurrent open -Work objects are permitted. Core still performs constant-time endpoint, -binding, and plan lifecycle checks and publishes all lifetime reservations -before calling a provider. For `host_wait`, static device and token validation -happens at bind time, while capture detection, event recording, and the actual -host wait happen on every submission. - -When a provider reports `FAILED`, Core derives and publishes the failure -diagnostic before committing the terminal Work state. If an asynchronous -`BaseException` interrupts that diagnostic lookup, Work remains nonterminal and -the next observation retries it; a terminal state can therefore never make a -missing diagnostic permanently unobservable. - -An ordinary `plan.bind(...)` binding remains fully immutable: operation, -indices, notification, CUDA ordering mode, stream, and ready event cannot be -overridden by `submit()`. The narrower `IndexSelection.bind()` form above owns -an immutable operation and index snapshot but intentionally accepts a -per-submission notification. Its provider path distinguishes a missing payload -from a replacement payload and cannot retain a stale tag. - -Providers receive an opaque `reuse_token`. Its identity is stable across all -submissions from one binding, allowing a provider to cache immutable native -lowering. An `IndexSelection` allocates its shared geometry token lazily on the -first selection-backed `BoundTransfer.submit()`: direct one-shot submissions -carry no token, and a binding used only to create a persistent request slot -allocates none. Selection-backed bindings also avoid a redundant private -binding token. The full-selection, notification-free, `caller_ready` -convenience forms `endpoint.read(plan)` and `endpoint.write(plan)` use stable -tokens per plan and operation. Tokens are hints only and never define transfer -semantics. - -### Persistent request slots - -A provider can advertise `Capabilities.persistent_request_slots` to expose the -lowest-allocation single-flight form of a bound transfer. A slot is created once -and owns one stable provider request; starting a generation does not allocate a -public `Work` or modify Core's resource sets: - -```python -bound = plan.bind( - _transfer.TransferOp.WRITE, - local_indices=[7, 2], - remote_indices=[19, 4], -) -slot = bound.create_slot() - -for _ in range(iterations): - slot.start_and_wait() - slot.recycle() - -slot.close() -bound.close() -``` - -For a blocking full-cycle loop, `execute()` combines start, completion, and -successful logical recycle under the slot's serialization boundary: - -```python -for _ in range(iterations): - completed_generation = slot.execute() -``` - -For queue-depth loops, `start_and_test()` exposes the initial post result with -zero additional status observations when a fused provider supports that exact -contract, while preserving Core ownership. It returns `True` with the slot logically -idle on immediate success, or `False` with an active generation. Poll pending -members with `test_and_recycle()`; it likewise returns `True` only after clean -successful completion has made the slot repostable. Failure and cancellation -are raised and remain owned for explicit recovery. Request-local notification -slots require `notification=` on `start_and_test`, exactly as on `start` and -`execute`. - -It returns the completed generation number. On success, `active` is false, -`state` is `None`, and the next generation or `close()` may proceed immediately. -Failure, cancellation, and timeout still leave the generation active, exactly -like `start_and_wait()`, so the caller can inspect or cancel it and must recycle -it explicitly after terminal state. The successful idle publication is atomic -with terminal observation under `SERIALIZED` and `MULTIPLE`; `SINGLE` provides -the same guarantee through its owner-thread contract, and -`CALLER_SERIALIZED` through the caller's whole-resource-graph exclusion. - -Providers may advertise -`Capabilities.deferred_successful_request_slot_recycle`. Core then omits the -provider `recycle()` call after clean `COMPLETED` in `execute()`, -`start_and_test()`, or `test_and_recycle()`. The provider retains that -observable completion receipt and consumes it inside the next `start()`, -`start_and_poll()`, or `close()` boundary. Explicit -`start()`/`wait()`/`start_and_wait()` behavior is unchanged and still requires -explicit `recycle()`. Providers without the capability remain source-compatible: -Core performs provider recycle before publishing execute success. - -Explicit `recycle()` completes provider recycle before making `active` false, -the sole public idle commit. Private cleanup follows, with `state = None` last -as a clean-idle marker. Public state and error ignore private terminal -diagnostics while idle; if asynchronous interruption stops cleanup, the next -begin observes the marker and idempotently repairs it before publishing another -generation. - -The exact `timeout=None` + lock-free + fused-start + deferred-success case is a -straight-line Core specialization. It applies to `SINGLE` and correctly used -`CALLER_SERIALIZED` endpoints and performs the same closed/creation, -conditional owner-thread, active-generation, CUDA-ordering, -generation-publication, and ambiguous-launch checks without crossing the -generic helper frames. Immediate clean completion publishes idle directly. -Because `_active` is its sole Core publication commit and the provider owns -durable state after dispatch entry, this exact synchronous path needs no -transient Core PENDING or success-side state-clear store; pre-entry interruption -becomes a retained Core-only failure. Nonterminal and exceptional results rejoin -the ordinary bounded/observed recovery paths. This is an internal latency -optimization only and does not change the public or provider contract. - -`TransferSlot` is intentionally mutable and single-flight. Its `state` is -`None` while idle. `start()` creates one owned generation, increments -`generation`, and changes `active` to true. On the explicit lifecycle, a second -start is rejected until the first generation reaches `COMPLETED`, `FAILED`, or -`CANCELLED` and the caller invokes `recycle()`. `active` remains true in terminal -state until that recycle boundary. Every `wait()` and `start_and_wait()` pins -the generation it began observing; if another thread recycles and restarts the -mutable slot before the waiter makes its next observation, that waiter raises -`BusyError` instead of reporting the replacement generation. A slot can close -only while idle. - -The ordinary binding and `TransferSlot` remain fully immutable: operation, -indices, attached notification, stream/event, and CUDA ordering cannot change -between generations. With a bound stream in `host_wait`, each slot owns one -reusable, lazily materialized CUDA event. Core caches its record/synchronize -callables and the stream capture query; every `start()` still checks capture, -records that same event after prior producer work, and host-waits before entering -the provider. An explicitly supplied ready event keeps its cached synchronize -callable. No readiness event is created in the steady generation path. - -When a framework reuses one exact operation/descriptor/index selection but -changes its completion tag, it may instead call -`binding.create_slot(request_local_notifications=True)`. The returned -`RequestLocalNotificationSlot` requires `notification=` on every `start`, -`start_and_wait`, or `execute`; `None` explicitly clears the native tag. This -mode is separately capability-gated so the fixed `TransferSlot.execute()` path -does not gain a request-local branch. The operation, selected indices, CUDA -ordering, and all memory ownership remain immutable. - -`start_and_wait(timeout)` is the latency-oriented equivalent of `start()` plus -`wait(timeout)`. A provider may additionally advertise -`Capabilities.fused_request_slot_start_poll` and implement the bounded -`BackendRequestSlot.start_and_poll(max_polls=..., timeout_ns=...)` SPI. Core -then keeps native post and up to 64 initial status observations inside one -provider boundary. The same SPI accepts `max_polls=0` for post-only execution: -an overriding fused provider returns the post-derived state, while the default -source composition can only return its cached `state`. A provider may independently advertise -`Capabilities.fused_request_slot_poll` and implement -`BackendRequestSlot.poll_bounded(max_polls=..., timeout_ns=..., -timeout_check_interval=...)`. If the request -is still pending or running, Core continues in chunks of at most 4,096 native -observations and a 100-microsecond soft residency target. The default interval -checks the timeout after observation 1 and then every 32 observations; one -provider call is non-preemptible and the sampled target can overshoot. Every -chunk reacquires the slot lock, verifies the pinned generation, invokes the -provider, publishes monotonic state, and releases the lock. This removes -per-observation Python crossings and most clock-read tax while retaining -count-bounded cancellation and signal checkpoints between chunks. Providers -without poll-only fusion retain the latency-oriented single-poll spin path; -providers without start fusion may still accelerate polling after ordinary -`start()`. - -At slot adoption, a provider may pass -`adopt_slot(wrapper, execution_dispatch=dispatch)`. Core strongly retains the -dispatch object and validates and caches its `start`, `poll_state`, `cancel`, -and `recycle` callables, plus `start_and_poll` and `poll_bounded` when their -capabilities are advertised. The wrapper remains authoritative for `error`, -`completion_event`, and `close`. Omitting `execution_dispatch` caches the -wrapper methods exactly as before, so existing providers remain source -compatible. The hook lets a native object own latency-sensitive execution -state without weakening Core's generation, timeout, CUDA-ordering, or lifetime -authority. - -For the deferred-success, fused, `SINGLE` `execute()` specialization, adoption -caches the eligible fused callable itself and whether CUDA readiness is already -caller-owned. The steady path branches on that callable before loading the -generic lock or dispatch table; serialized, timed, non-fused, and non-deferred -slots continue through the general path. - -Request-slot operations are interruption-safe but not signal-handler-reentrant: -after a Python signal interrupts an operation, let that frame unwind before -calling the same slot again. A subsequent ordinary call recovers the durable -state. This avoids charging every hot dispatch for a recursive-entry guard that -cannot make arbitrary signal-handler use safe. - -The standalone `wait()` operation retains cooperative yield/backoff for callers -that prefer lower host-CPU pressure; both operations have identical state, -timeout, cancellation, and lifetime semantics. - -Core passes the provider timeout as integer nanoseconds and saturates extremely -large finite public timeouts at signed 64-bit maximum for that bounded call; -the subsequent Core deadline remains an arbitrary-precision Python integer. -The private, introspectable policy constants are -`_FUSED_REQUEST_SLOT_POLL_MAX_POLLS` and -`_FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS`, and -`_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL`; benchmarks should read them rather than -duplicate the current 4,096/100-microsecond policy. - -The timeout is a total host-side budget from entry to `start_and_wait`. It -cannot preempt one blocking CUDA-readiness or provider call. A timeout never -recycles implicitly: the generation remains active and cancellable, and the -caller must observe a terminal state before explicitly invoking `recycle()`. -Every generation, fused or not, repeats `host_wait` readiness before Core -publishes the active generation and enters the provider. - -`completion_event` is exposed only when the backend advertises -`cuda_completion_event`, and only while a generation is active. The event is -generation-local: enqueue any device wait before `recycle()`, because a later -generation may re-record or replace the provider event. - -The provider slot is adopted only after Core has published the slot in the -complete endpoint/binding/plan lifetime graph and before the provider may create -native state. Once `start()` crosses into the provider, Core conservatively owns -that generation even if an asynchronous exception crosses the call boundary. -Providers must convert failures after possible native launch into terminal slot -state. The retained generation can then be inspected, cancelled where possible, -and recycled without losing native ownership. - -Slot creation publishes the Endpoint recovery edge before the binding blocker, -so an interruption cannot leave an inaccessible child that permanently blocks -its parent. Close uses closure as a retryable commit marker: it retains at least -one parent edge until that marker is visible, and a repeated `close()` finishes -any interrupted Core-only unlink/cleanup epilogue. Generation publication and -terminal failure/cancellation publication use the same commit-last rule. A -pre-provider interruption can therefore be recycled as a Core-only failed -generation without issuing provider `recycle()` against an idle native slot; -an interrupted terminal-error lookup is retried on the next observation. - -This opt-in path is for repeated immutable transfers. `BoundTransfer.submit()` -remains the general concurrent path and continues to return an independent -`Work` for every submission. A TransferSlot never makes `Work` reusable. - -## Work and ownership - -`Work.test()` is request-local and returns whether the request is terminal. -`Work.state` distinguishes completed, failed, and cancelled work; `Work.error` -retains the failure. `wait(timeout)` raises typed timeout, cancellation, or -backend errors. Cancellation is best effort and does not roll back remote writes. - -Core retains the plan, registrations, peer, ordering tokens, and notification -payload until Work is closed. Closing a live or referenced resource raises -`BusyError`. Close operations are idempotent after success. Required teardown is: - -```text -Work -> BoundTransfer (when used) -> TransferPlan -> Peer / Registration -> Endpoint -TransferSlot -> BoundTransfer -> TransferPlan -> Peer / Registration -> Endpoint -``` - -`Endpoint.open_bound_transfers` is the recovery snapshot for bindings. A -binding cannot close while it has an open Work or TransferSlot, a plan cannot -close while it has an open binding, and an endpoint cannot close until all -bindings are closed. `Endpoint.open_transfer_slots` exposes idle, in-flight, -terminal, and cleanup-only slots for recovery after interruption. - -One-shot READ/WRITE owns an internal plan and releases it when Work closes. -The provider SPI uses an explicit adoption handshake for data transfers and -standalone notifications: Core publishes the Work and its complete lifetime -graph first, then the provider must adopt a `BackendWork` before launching -native activity. A provider may raise before adoption only when no native work -remains; after adoption, launch errors must become terminal Work failures. This -closes the interruption window between native submission and Python -bookkeeping. `Endpoint.open_works` provides a recovery snapshot if an -asynchronous exception lands after submission but before caller code stores the -returned Work. - -Code that must close that final caller-side adoption window may pass a unique -opaque `recovery_token=` to any Work-producing operation. Core stores the -object by identity before publishing the Work and exposes it unchanged through -`Work.recovery_token`. Recovery is deliberately a cold-path scan: - -```python -attempt = object() -try: - work = plan.submit_indices( - "write", - local_indices=local_indices, - remote_indices=remote_indices, - recovery_token=attempt, - ) -except BaseException: - matches = [ - candidate - for candidate in endpoint.open_works - if candidate.recovery_token is attempt - ] - # zero: Core retained no provider activity; one: the lost Work; - # more than one: an ambiguous token-uniqueness error. - raise -``` - -Core never hashes the token, compares it, coerces it to Boolean, invokes it, -serializes it, or passes it to the provider. It is held strongly until -`Work.close()` removes every lifetime edge, then cleared as the close -epilogue's last commit. Use `is`, never `==`, and do not use the default `None` -for recovery matching. This is distinct from provider-facing `reuse_token`: -the reuse token identifies a cached native request shape; the recovery token -identifies one caller submission attempt. - -Registration, peer import, and plan preparation use the same ownership rule: -Core publishes a resource wrapper, the provider adopts a retry-safe transaction -record before its first native mutation, and Core rolls it back on failure. If -rollback itself fails, the wrapper is cleanup-only and remains recoverable via -`Endpoint.open_registrations`, `open_peers`, or `open_plans`; `close()` retries -the provider cleanup. - -Endpoint construction is the one boundary where Core cannot first publish a -recovery object: Python has not returned the `Endpoint` to its caller yet. -Provider factories must therefore catch `BaseException` and self-roll back all -native state created before they return. This is a provider obligation, not an -exactly-once guarantee that Core can enforce around arbitrary third-party -factory code. - -## Progress and threading - -`progress_mode` is `manual` or `background`. Background progress is provider -owned; Core does not create a hidden thread. `thread_mode` is `single`, -`caller_serialized`, `serialized`, or `multiple`: - -- `single` is confined to the thread that created the Endpoint and has no Core - locks; -- `caller_serialized` may migrate between threads, but the application must - prevent overlap and reentry for the entire Endpoint resource graph, including - Work/slot observation, progress, cancellation, snapshots, waits, and teardown; -- `serialized` lets calls migrate and makes Core provide that exclusion, - including provider entry; -- `multiple` permits concurrent operations when the provider advertises that - capability. Core publishes Work/lifetime reservations before releasing its - lifecycle lock, so data submission and status, standalone notification send, - destructive notification receive, and manual progress can overlap when their - backend resources are independent. Destructive notification consumers remain - serialized with each other. Endpoint close returns `BusyError` while an - admitted provider call is active and never tears down its target concurrently. - Providers still serialize private destructive state that cannot safely join - that overlap. - -Provider-owned background progress threads are outside the caller boundary. -An application must keep its external exclusion for the complete public call, -including a blocking `host_wait` or `await Work` iteration. Selecting -`caller_serialized` without that whole-graph contract is unsafe. - -Manual progress is endpoint-local. Every participating process must call -`Endpoint.progress()` regularly, including an otherwise idle remote responder; -waiting on one endpoint cannot drive a different process's transport engine. - -## Device ordering - -READ and WRITE expose two explicit CUDA producer-ordering modes: - -- `caller_ready` (default): the caller guarantees that selected local memory is - ready before submission. No CUDA API is called. -- `host_wait`: pass exactly one real `torch.cuda.Stream` or recorded - `torch.cuda.Event`. Core validates that every selected local Region is on one - explicitly indexed CUDA device, records one event when a stream is supplied, - then waits for that event on the host before posting the provider request. - -The wait happens outside Core's state and provider locks. It never inserts a -device-wide synchronization, and providers must not add another implicit wait. -`host_wait` protects local memory only: readiness of a remote READ source still -belongs to the remote application's control protocol. Core checks capture on a -supplied stream before recording its event. A supplied ready event goes directly -to its authoritative `synchronize()` operation and does not query the unrelated -current stream's capture state; an event CUDA cannot synchronize still fails -through CUDA. -Native stream-attached, graph-safe transport submission would require a future, -separately advertised mode. - -Terminal Work completion means the local allocation can be reused. A provider -may additionally advertise `cuda_completion_event`, in which case -`Work.completion_event` expresses that terminal point to CUDA. The initial NIXL -provider does not advertise such an event. - -## Notifications - -Capabilities distinguish two contracts that native integrations use -differently: - -- `attached_notifications`: a notification attached to READ/WRITE is - transfer-ordered and cannot become observable before the associated remote - writes are visible. -- `standalone_notifications`: `send_notification` has the terminal meaning in - `standalone_notification_completion`. NIXL advertises `local_accepted`: a - completed Work means the native API accepted/copied the payload, not that the - remote application observed it. Remote delivery still requires transport - progress and receiver polling. - -A backend may additionally advertise `synchronous_notification_send` and -implement `Endpoint.send_notification_sync(peer, payload)`. This additive fast -path returns `None` and creates no Core or provider `Work`; return still means -only `standalone_notification_completion`. For NIXL that is `local_accepted`, -not remote observation. An exception after provider entry can therefore be -outcome-ambiguous. Reliable control protocols must retain sequence IDs, -acknowledgement, deduplication, and retry. Callers that need inspectable local -ownership after interruption should keep using `send_notification()`. -In `multiple` mode Core uses a short peer-lifetime reservation and releases its -endpoint lifecycle lock before provider entry, so independent synchronous sends -retain provider/native concurrency. The reservation is a pooled strong token in -the send frame plus a close-visible weak identity: steady state allocates no new -token, and an interrupted cleanup leaves a dead identity that `Peer.close()` -prunes rather than a permanent false blocker. `Peer.close()` reports -`BusyError` only while a live reservation exists. Providers rely on this Core -pin and should not duplicate it on their hot SPI path. `serialized` continues -to exclude the complete call, while `single` and `caller_serialized` retain -their caller contracts. - -The asynchronous `send_notification()` path applies the same lock-breadth rule -without another token allocation: its required `Work` is published with the -peer and payload lifetime anchors before provider entry. Progress and receive -use pooled close-visible call reservations, allocated only when concurrency -exceeds the warmed high-water mark. Receive calls take a dedicated consumer -lock because dequeue is destructive; that lock does not exclude data submit, -independent Work status, asynchronous sends, or provider-safe progress. - -Received notifications carry either no source identity or the complete -`(endpoint_id, incarnation)` pair. Consumers must reject stale or unknown -incarnations rather than routing by a reusable display name alone. - -`poll_notifications()` is an at-most-once receive boundary. The current SPI -allows a provider to dequeue before returning the batch, so an asynchronous -exception or allocation failure after dequeue can lose that batch. Protocols -that require reliable control messages must include an application sequence ID, -acknowledgement, deduplication, and retransmission. A future peek/token/ack SPI -would be required to provide a stronger receive transaction. - -High-rate consumers should use -`poll_notification_batches(max_items=None)`. `max_items` counts payloads across -all returned batches; zero returns without entering the provider. Each -`NotificationBatch` contains a nonempty immutable `tuple[bytes, ...]` and either -a complete source identity or no identity. Repeated batches for one source are -legal and preserve that source's observation order. Batch boundaries have no -wire meaning, and no total order is promised between different sources. - -Providers may advertise `grouped_notification_receive` when they retain their -native source grouping instead of first constructing one provider object per -payload. The public operation remains available without that capability through -a scalar compatibility implementation. It groups only adjacent notifications -with the same complete identity and keeps anonymous notifications separate. -Payload tuples that are already exact immutable bytes are retained by identity; -mutable inputs are snapshotted. Returned-value validation and public object -construction run after Core releases its endpoint state lock, while the -destructive provider call itself remains serialized with close and mutation. - -`max_items` is an egress bound, not automatically a native-ingress memory bound. -A provider whose native API drains a complete transport batch before returning -must separately cap or document that temporary materialization. The current -NIXL provider bounds retained Python count and bytes and fails closed on -malformed/overflowed ingress, but the underlying NIXL `getNotifs` API does not -yet offer a bounded dequeue. - -## Metadata and security boundary - -`EndpointMetadata` uses bounded canonical JSON with base64 opaque provider -fields, a `{major, minor}` wire version, required-feature names, optional opaque -extensions, stable endpoint ID, and restart incarnation. Unknown outer fields, -missing or duplicate fields, unknown required features, malformed values, and -unsupported major versions are rejected. A higher minor version is accepted -only through understood required features and the extension container. Encoded -metadata, provider payloads, descriptors, registration counts, notification -payloads, and identity strings have explicit limits. Opaque capability-bearing -provider data is redacted from object representations. No pickle is used by -Core. - -Callers with a stable logical identity may pass `endpoint_id=`. Core creates a -fresh incarnation for each Endpoint construction and rejects replacement -metadata until the old Peer is explicitly closed. - -When the control plane already knows a peer identity, use -`import_peer(expected_identity=...)`, or the individual `expected_endpoint_id`, -`expected_name`, and `expected_incarnation` checks. Core performs these checks -before the provider can mutate native peer state. - -`export_metadata(registrations=[...])` is a capability grant for exactly the -selected registrations. Providers must emit partial transport metadata or reject -the operation; silently exposing omitted registrations is not allowed. -In this initial contract each imported Peer is an immutable grant snapshot. -Extending or revoking that grant requires closing dependent plans and the Peer, -then importing a replacement snapshot. Incremental `Peer.update` semantics are -left open because merge, revocation, and in-flight-work ordering must be defined -together. - -Metadata authenticity, confidentiality, replay protection, and application-level -authorization remain the control plane's responsibility. Metadata is a bearer -capability: do not send it over an unauthenticated channel or log its opaque -fields. Core's bounds and strict parser are defense in depth, not a secure -transport envelope. - -The exporting endpoint cannot observe remote plans or work. The application -must therefore hold a registration lease until every peer has quiesced and -released work that may reference it, then communicate revocation before local -deregistration. Core's local lifetime anchors do not implement a distributed -lease protocol. - -## Initial limitations - -- private Python API and SPI; no stable ABI; -- CPU reference backend only for conformance; -- the NIXL provider has two-node GB200/UCX correctness and large-transfer - performance coverage, but not full framework end-to-end GPU coverage; -- no atomic model activation, remote lease protocol, or rollback; -- no FILE/OBJ tier, distributed-checkpoint durability, or NIXL-EP token path; -- `host_wait` is a host-side gate, not a native GPUNETIO stream submission; -- no CUDA Graph or compiler integration and no CUDA completion event from NIXL; -- no promise that cancellation prevents partial remote visibility. -- full-catalog overlap checks conservatively reject compact strided regions - whose bounding envelopes interleave; explicit indexed selections are checked - exactly without expanding the whole catalog. +assert target.tolist() == [4, 5, 6, 7, 0, 1, 2, 3] +a.close() +b.close() +``` + +For NIXL, construct `Endpoint(nixl.torch_transfer.NixlBackend(native_agent))`. +The provider stays in NIXL. `MemoryBackend` is CPU-only, synchronous and in-process. + +## Contract + +- `register_tensor` retains a contiguous tensor. Raw `register` takes rows + `(address, bytes, device)` and a strong owner; the caller guarantees validity. +- `prepare` caches a descriptor catalog; optional `(stride, count)` fields keep + strided layouts compact. The provider validates indices/ranges/paired lengths. +- `submit` selects equally sized lists of indices. WRITE reads local memory; + READ writes local memory. Destination spans must not overlap. +- All methods serialize host access with one endpoint lock. GPU ordering is + separate: `caller_ready` trusts the caller; `host_wait` synchronizes the supplied + CUDA event before posting. No CUDA graph or stream-ordered completion support. +- `Work.poll()` returns `pending`, `done` or `error`. `wait()` timeout does not + cancel. Close succeeds only after `done`; pending/error work retains ownership. +- Close work, catalogs, peers, then registrations (or call endpoint.close). + Registration teardown conservatively waits for ALL endpoint catalogs to close. +- Stop remote access before teardown. Applications own rendezvous, authentication, + allocation lifetime, leases and acknowledgement protocols. Standalone `notify` + means local acceptance, not remote acknowledgement. Bytes remain opaque. + +## Deliberately deferred + +Active cancellation, failure recovery, revocation, metadata schemas, capability +negotiation, native dispatch helpers, slots, graph capture and performance parity. +No asynchronous Python interruption guarantee. Native GPU/end-to-end framework +validation is required before expanding this prototype or claiming speed parity. diff --git a/torch/distributed/_transfer/__init__.py b/torch/distributed/_transfer/__init__.py index 9c181f25dea2f..d5e117a599eaa 100644 --- a/torch/distributed/_transfer/__init__.py +++ b/torch/distributed/_transfer/__init__.py @@ -1,123 +1,9 @@ -"""Experimental backend-neutral endpoint transfer API. +"""Experimental, backend-independent point-to-point memory transfers. -This module is private while its ownership, ordering, and provider contracts are -validated with real inference and training consumers. +This is a review prototype, not a stable or production-supported API. """ -from ._api import ( - BoundTransfer, - Endpoint, - IndexSelection, - Peer, - RawSpan, - RequestLocalNotificationSlot, - Region, - Registration, - TransferPlan, - TransferSlot, - Work, -) -from ._backend import ( - BACKEND_FACTORY_API_VERSION, - Backend, - BackendFactory, - BackendFactoryV2, - BackendIndexSelection, - BackendRequestSlot, - BackendWork, - ImportedRawMemory, - ImportedRegistration, - RegisteredMemory, - register_backend, - unregister_backend, -) -from ._errors import ( - BackendError, - BackendFailureError, - BackendMismatchError, - BackendNotFoundError, - BusyError, - ClosedError, - InvalidMetadataError, - InvalidRegionError, - TransferCancelledError, - TransferError, - TransferTimeoutError, - UnsupportedError, -) -from ._reference import _register_reference_backend -from ._types import ( - BackendNotification, - BackendNotificationBatch, - BackendRegion, - Capabilities, - CudaOrderingMode, - EndpointIdentity, - EndpointInfo, - EndpointMetadata, - Notification, - NotificationBatch, - NotificationCompletion, - ProgressMode, - RegistrationMetadata, - ThreadMode, - TransferOp, - WorkState, -) +from ._api import Catalog, Endpoint, Registration, Work +from ._backend import Backend - -_register_reference_backend() - -__all__ = [ - "BACKEND_FACTORY_API_VERSION", - "Backend", - "BackendError", - "BackendFactory", - "BackendFactoryV2", - "BackendIndexSelection", - "BackendFailureError", - "BackendMismatchError", - "BackendNotFoundError", - "BackendNotification", - "BackendNotificationBatch", - "BackendRegion", - "BackendRequestSlot", - "BackendWork", - "BoundTransfer", - "BusyError", - "Capabilities", - "ClosedError", - "CudaOrderingMode", - "Endpoint", - "EndpointIdentity", - "EndpointInfo", - "EndpointMetadata", - "ImportedRawMemory", - "ImportedRegistration", - "IndexSelection", - "InvalidMetadataError", - "InvalidRegionError", - "Notification", - "NotificationBatch", - "NotificationCompletion", - "Peer", - "ProgressMode", - "RawSpan", - "RequestLocalNotificationSlot", - "Region", - "RegisteredMemory", - "Registration", - "RegistrationMetadata", - "ThreadMode", - "TransferCancelledError", - "TransferError", - "TransferOp", - "TransferPlan", - "TransferSlot", - "TransferTimeoutError", - "UnsupportedError", - "Work", - "WorkState", - "register_backend", - "unregister_backend", -] +__all__ = ["Backend", "Catalog", "Endpoint", "Registration", "Work"] diff --git a/torch/distributed/_transfer/_api.py b/torch/distributed/_transfer/_api.py index 06fde259e7a85..1e10d2688f07e 100644 --- a/torch/distributed/_transfer/_api.py +++ b/torch/distributed/_transfer/_api.py @@ -1,8283 +1,259 @@ from __future__ import annotations -import asyncio -import ctypes -import functools -import itertools -import math -import operator -import sys import threading import time -import uuid -import weakref -from bisect import bisect_right -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Any, cast, ParamSpec, TypeVar +from typing import Any import torch -from typing_extensions import Self -try: - from torch._C._distributed_transfer import ( # type: ignore[attr-defined] - _snapshot_and_analyze_int32_indices as _compiled_index_analyzer, - ) -except ImportError: - # Source-tree Python can run against an older installed torch._C while this - # private API is under development. Mixed/general sequences also never - # depend on the compiled helper. - _compiled_index_analyzer = None +from ._backend import Backend -from ._backend import ( - Backend, - BackendIndexSelection, - BackendRequestSlot, - BackendWork, - ImportedRawMemory, - ImportedRegistration, - RegisteredMemory, - _get_backend_registration, - _make_backend_index_selection, -) -from ._errors import ( - BackendFailureError, - BackendMismatchError, - BusyError, - ClosedError, - InvalidMetadataError, - InvalidRegionError, - TransferCancelledError, - TransferError, - TransferTimeoutError, - UnsupportedError, -) -from ._types import ( - BackendNotification, - BackendNotificationBatch, - BackendRegion, - Capabilities, - CudaOrderingMode, - EndpointIdentity, - EndpointInfo, - EndpointMetadata, - Notification, - NotificationBatch, - ProgressMode, - RegistrationMetadata, - ThreadMode, - TransferOp, - WorkState, - _notification_batch_from_backend, -) +# Dropping a Python Work reference must not free memory still used by DMA. +# Explicit close removes the root; there is deliberately no GC-driven abort. +_live_endpoints: set[Endpoint] = set() -__all__ = [ - "BoundTransfer", - "Endpoint", - "IndexSelection", - "Peer", - "RawSpan", - "RequestLocalNotificationSlot", - "Region", - "Registration", - "TransferPlan", - "TransferSlot", - "Work", -] -_TERMINAL_STATES = frozenset( - {WorkState.COMPLETED, WorkState.FAILED, WorkState.CANCELLED} -) -_CONSTRUCTION_TOKEN = object() -_UNADOPTED_HANDLE = object() -_UNSET_NOTIFICATION = object() -_MAX_NOTIFICATION_BYTES = 1024 * 1024 -_FUSED_REQUEST_SLOT_MAX_POLLS = 64 -_MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS = (1 << 63) - 1 -_FUSED_REQUEST_SLOT_POLL_MAX_POLLS = 4096 -_FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS = 100_000 -_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL = 32 -_GET_IDENT = threading.get_ident -_WORK_COMPLETED = WorkState.COMPLETED -_WORK_RUNNING = WorkState.RUNNING -_NATIVE_INT32_FORMATS = frozenset( - {"i", "@i", "=i", "i"} -) -_INT32_MAX = (1 << 31) - 1 -_INDEX_SELECTION_TORCH_THRESHOLD = 128 -_P = ParamSpec("_P") -_R = TypeVar("_R") - - -class _BackendConstructionReceipt: - """Core ownership record independent of an Endpoint object's lifetime.""" +class Endpoint: + """One provider instance and its locally owned resources. - __slots__ = ( - "_adoption_count", - "_backend", - "_closed", - "_lock", - "_poison_message", - "_published", - "_sealed", - "backend_name", - "endpoint_id", - ) + Control-plane exchange, remote allocation lifetime and disjoint destination + writes belong to the application. Never resize/reallocate registered memory. + Host calls are serialized; this does NOT order GPU streams. + """ - def __init__(self, backend_name: str, endpoint_id: str) -> None: - self.backend_name = backend_name - self.endpoint_id = endpoint_id + def __init__(self, backend: Backend): + self._backend = backend self._lock = threading.RLock() - self._backend: Backend | None = None - self._adoption_count = 0 - self._sealed = False - self._published = False + self._registrations: set[Registration] = set() + self._catalogs: set[Catalog] = set() + self._works: set[Work] = set() + self._peers: set[str] = set() self._closed = False - self._poison_message: str | None = None - - @property - def backend(self) -> Backend | None: - with self._lock: - return self._backend - @property - def published(self) -> bool: - with self._lock: - return self._published - - @property - def closed(self) -> bool: - with self._lock: - return self._closed + def _check_open(self): + if self._closed: + raise RuntimeError("endpoint is closed") - def adopt(self, backend: Backend) -> None: - if not isinstance(backend, Backend): - message = "backend factory must adopt a Backend instance" - with self._lock: - if self._poison_message is None: - self._poison_message = message - raise TypeError(message) + def register(self, regions: Any, *, kind: str, owner: Any) -> Registration: + """Register raw regions; owner must keep ALL backing allocations alive.""" + if kind not in ("host", "cuda") or owner is None: + raise ValueError("expected host/cuda memory and a strong owner") with self._lock: - if self._sealed: - message = "backend factory attempted adoption after its call returned" - if self._poison_message is None: - self._poison_message = message - raise BackendFailureError(message) - if self._backend is not None or self._adoption_count: - message = "backend factory must call adopt_backend exactly once" - if self._poison_message is None: - self._poison_message = message - raise BackendFailureError(message) - # Ownership commits before the diagnostic count. An interruption - # between these stores still leaves the native owner closeable. - self._backend = backend - self._adoption_count = 1 + self._check_open() + reg = Registration(self, self._backend.register(regions, kind), owner) + self._registrations.add(reg) + return reg + + def register_tensor(self, tensor: torch.Tensor) -> Registration: + """Register the byte span of a contiguous CPU or CUDA tensor.""" + if tensor.device.type not in ("cpu", "cuda") or not tensor.is_contiguous(): + raise ValueError("expected a contiguous CPU or CUDA tensor") + if tensor.numel() == 0: + raise ValueError("empty registration") + regions = [ + ( + tensor.data_ptr(), + tensor.numel() * tensor.element_size(), + tensor.device.index or 0, + ) + ] + return self.register( + regions, kind="cuda" if tensor.is_cuda else "host", owner=tensor + ) - def seal_failed_call(self) -> None: + def metadata(self) -> bytes: + """Export provider metadata AFTER registration; exchange out of band.""" with self._lock: - self._sealed = True + self._check_open() + return self._backend.metadata() - def seal_successful_call(self, result: object) -> Backend: + def connect(self, metadata: bytes) -> str: with self._lock: - self._sealed = True - if self._poison_message is not None: - raise BackendFailureError(self._poison_message) - if result is not None: - raise TypeError("version-2 backend factory must return None") - if self._adoption_count != 1 or self._backend is None: - raise BackendFailureError( - "version-2 backend factory must call adopt_backend exactly once" - ) - return self._backend + self._check_open() + peer = self._backend.connect(metadata) + self._peers.add(peer) + return peer - def mark_published(self) -> None: + def disconnect(self, peer: str) -> None: with self._lock: - if self._closed or self._backend is None: - raise BackendFailureError("cannot publish a closed backend receipt") - self._published = True + if any(c.peer == peer for c in self._catalogs): + raise RuntimeError("close peer catalogs before disconnecting") + if peer in self._peers: + self._backend.disconnect(peer) + self._peers.remove(peer) - def close(self, close_backend: Callable[[], object] | None = None) -> None: - """Close once; an interrupted/failed close remains retryable.""" + def prepare(self, regions: Any, *, kind: str, peer: str | None = None) -> Catalog: + """Prepare a reusable descriptor catalog, local if peer is None. + The provider checks registration coverage. The application supplies + remote addresses/layout alongside metadata through its control plane. + """ with self._lock: - if self._closed: - return - backend = self._backend - if backend is None: - self._closed = True - return - callback = backend.close if close_backend is None else close_backend - try: - result = callback() - if result is not None: - raise BackendFailureError("backend close must return None") - except BaseException: - raise - self._closed = True - - -_ORPHANED_BACKEND_RECEIPTS: dict[tuple[str, str], set[_BackendConstructionReceipt]] = {} -_ORPHANED_BACKEND_RECEIPTS_LOCK = threading.Lock() - - -def _retain_orphaned_backend_receipt(receipt: _BackendConstructionReceipt) -> None: - key = (receipt.backend_name, receipt.endpoint_id) - with _ORPHANED_BACKEND_RECEIPTS_LOCK: - _ORPHANED_BACKEND_RECEIPTS.setdefault(key, set()).add(receipt) - - -def _forget_orphaned_backend_receipt(receipt: _BackendConstructionReceipt) -> None: - key = (receipt.backend_name, receipt.endpoint_id) - with _ORPHANED_BACKEND_RECEIPTS_LOCK: - receipts = _ORPHANED_BACKEND_RECEIPTS.get(key) - if receipts is None: - return - receipts.discard(receipt) - if not receipts: - _ORPHANED_BACKEND_RECEIPTS.pop(key, None) - - -def _finalize_backend_receipt(receipt: _BackendConstructionReceipt) -> None: - """Best-effort GC cleanup without retaining the Endpoint itself.""" - - if receipt.backend is None or receipt.closed: - return - _retain_orphaned_backend_receipt(receipt) - try: - # No Endpoint remains to supply its SERIALIZED lock. The receipt lock - # still serializes finalizer cleanup and all later orphan retries. - receipt.close() - except BaseException: - return - _forget_orphaned_backend_receipt(receipt) - - -def _retry_orphaned_backend_receipts(backend_name: str, endpoint_id: str) -> None: - """Audit one provider's orphans; only an exact identity blocks creation.""" - - with _ORPHANED_BACKEND_RECEIPTS_LOCK: - snapshot = tuple( - (key, tuple(receipts)) - for key, receipts in _ORPHANED_BACKEND_RECEIPTS.items() - if key[0] == backend_name - ) - - exact_failures: list[tuple[_BackendConstructionReceipt, Exception]] = [] - for key, receipts in snapshot: - for receipt in receipts: - try: - receipt.close() - except BaseException as error: - # Process-control interruptions are not converted or hidden; - # the retained receipt makes the next attempt safe. - if not isinstance(error, Exception): - raise - if key[1] == endpoint_id: - exact_failures.append((receipt, error)) - else: - _forget_orphaned_backend_receipt(receipt) - - if exact_failures: - receipt, error = exact_failures[0] - origin = ( - "an abandoned published endpoint (possibly with live provider children)" - if receipt.published - else "an unpublished endpoint construction" - ) - failure = BackendFailureError( - f"cannot reuse endpoint identity {endpoint_id!r} for backend " - f"{backend_name!r}: cleanup of {origin} is still pending" - ) - add_note = getattr(failure, "add_note", None) - if callable(add_note): - add_note(f"orphan cleanup failure type: {type(error).__name__}") - raise failure from error - - -class _NoopLock: - """Context-manager lock for SINGLE or externally serialized resources.""" - - __slots__ = () - - def __enter__(self) -> _NoopLock: - return self - - def __exit__( - self, - exc_type: object, - exc_value: object, - traceback: object, - ) -> bool: - return False - - -_NOOP_LOCK = _NoopLock() - - -class _ReuseToken: - """Weak-referenceable identity token with no provider-visible semantics.""" - - __slots__ = ("__weakref__",) - - -class _SyncNotificationReservation: - """Pooled strong token with a close-visible weak identity.""" - - __slots__ = ("active", "reference", "__weakref__") - - def __init__(self) -> None: - self.active = False - self.reference = weakref.ref(self) - - -class _BackendCallReservation: - """Reusable MULTIPLE-call token whose weak identity is visible to close.""" - - __slots__ = ("active", "reference", "__weakref__") - - def __init__(self) -> None: - self.active = False - self.reference = weakref.ref(self) - - -@dataclass(slots=True) -class _LifetimeAnchor: - owner: object - storage: object | None = None - buffer: memoryview | None = None - allocation_address: int | None = None - allocation_nbytes: int | None = None - - -@dataclass(frozen=True, slots=True) -class _RemoteRawDescriptor: - address: int - nbytes: int - device: str - memory_type: str - - -def _state_locked(function: Callable[_P, _R]) -> Callable[_P, _R]: - """Serialize endpoint resource-graph mutations.""" - - @functools.wraps(function) - def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: - endpoint = cast(Any, args[0]) - state_lock = endpoint._state_lock - if state_lock is _NOOP_LOCK: - return function(*args, **kwargs) - with state_lock: - return function(*args, **kwargs) - - return wrapper - - -@dataclass(frozen=True, slots=True) -class RawSpan: - """Expert-only local address range with an explicit lifetime anchor.""" - - address: int - nbytes: int - device: str - memory_type: str - owner: object = field(repr=False) - - def __post_init__(self) -> None: - _validate_positive_int(self.address, "address") - _validate_positive_int(self.nbytes, "nbytes") - _validate_nonempty_str(self.device, "device", InvalidRegionError) - _validate_nonempty_str(self.memory_type, "memory_type", InvalidRegionError) - if self.owner is None: - raise InvalidRegionError("a local RawSpan requires a lifetime owner") - - -class Region: - """A bounded local or remote byte region. - - Regions are created by :meth:`Registration.region`, :meth:`Peer.region`, - or :meth:`Peer.unsafe_region`; users do not construct them directly. - """ - - __slots__ = ( - "_backend_descriptor", - "_device", - "_memory_type", - "_nbytes", - "_offset", - "_origin", - "_remote", - "_stride", - "_count", - ) + self._check_open() + if kind not in ("host", "cuda"): + raise ValueError("expected host/cuda memory") + if peer is not None and peer not in self._peers: + raise ValueError("connect the peer first") + catalog = Catalog(self, self._backend.prepare(regions, kind, peer), peer) + self._catalogs.add(catalog) + return catalog - def __init__( + def submit( self, - origin: Registration | Peer, - backend_descriptor: object, + operation: str, + local: Catalog, + local_indices: Any, + remote: Catalog, + remote_indices: Any, *, - offset: int, - nbytes: int, - device: str, - memory_type: str, - remote: bool, - stride: int | None, - count: int | None, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("Region objects are created by Registration or Peer") - self._origin = origin - self._backend_descriptor = backend_descriptor - self._offset = offset - self._nbytes = nbytes - self._device = device - self._memory_type = memory_type - self._remote = remote - self._stride = stride - self._count = count - - @property - def offset(self) -> int: - return self._offset - - @property - def nbytes(self) -> int: - """Number of bytes in each strided segment.""" - - return self._nbytes - - @property - def device(self) -> str: - return self._device - - @property - def memory_type(self) -> str: - return self._memory_type - - @property - def stride(self) -> int | None: - return self._stride - - @property - def count(self) -> int | None: - return self._count - - @property - def total_nbytes(self) -> int: - return self._nbytes * (self._count or 1) - - @property - def is_remote(self) -> bool: - return self._remote + notification: bytes = b"", + ordering: str = "caller_ready", + ready_event: Any = None, + ) -> Work: + """Submit READ (remote -> local) or WRITE (local -> remote). - def _to_backend(self) -> BackendRegion: - self._origin._check_open() - descriptor = self._backend_descriptor - if isinstance(self._origin, Registration): - descriptor = self._origin._handle - elif isinstance(descriptor, RegistrationMetadata): - descriptor = ImportedRegistration( - peer_handle=self._origin._handle, - registration_id=descriptor.registration_id, - name=descriptor.name, - nbytes=descriptor.nbytes, - device=descriptor.device, - memory_type=descriptor.memory_type, - descriptor=descriptor.backend_descriptor, - ) - elif isinstance(descriptor, _RemoteRawDescriptor): - descriptor = ImportedRawMemory( - peer_handle=self._origin._handle, - address=descriptor.address, - nbytes=descriptor.nbytes, - device=descriptor.device, - memory_type=descriptor.memory_type, + caller_ready: caller has completed source production and prior uses + of the local destination. host_wait: synchronize the supplied CUDA + event first. The remote owner must independently order remote memory. + Neither mode makes completion stream-ordered or graph-capturable. + """ + if operation not in ("READ", "WRITE"): + raise ValueError("expected READ or WRITE") + if ordering == "host_wait" and ready_event is not None: + ready_event.synchronize() + elif ordering != "caller_ready" or ready_event is not None: + raise ValueError("host_wait requires an event; caller_ready takes none") + with self._lock: + self._check_open() + if ( + local not in self._catalogs + or remote not in self._catalogs + or local.peer is not None + or remote.peer is None + ): + raise ValueError("expected live local/remote catalogs on this endpoint") + if len(local_indices) == 0 or len(local_indices) != len(remote_indices): + raise ValueError("expected equal, nonempty index selections") + handle = self._backend.create_work( + operation, + local._handle, + local_indices, + remote._handle, + remote_indices, + notification, ) - return BackendRegion( - owner=self._origin, - offset=self._offset, - nbytes=self._nbytes, - device=self._device, - memory_type=self._memory_type, - backend_descriptor=descriptor, - remote=self._remote, - stride=self._stride, - count=self._count, - ) - - -class Registration: - """A backend registration that retains its source allocation. - - Retention does not freeze mutable Storage metadata. The caller must keep - the allocation address and extent unchanged until this object is closed; - in particular, resizing while a plan or Work exists is unsafe. Core's - pointer/size checks are diagnostics and cannot close the mutation race. - """ - - __slots__ = ( - "_closed", - "_creation_failed", - "_endpoint", - "_handle", - "_memory", - ) - - def __init__( - self, - endpoint: Endpoint, - memory: RegisteredMemory, - handle: object, - *, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("Registration objects are created by Endpoint") - self._endpoint = endpoint - self._memory = memory - self._handle = handle - self._creation_failed = False - self._closed = False - - @property - def registration_id(self) -> str: - return self._memory.registration_id - - @property - def name(self) -> str: - return self._memory.name - - @property - def nbytes(self) -> int: - return self._memory.nbytes - - @property - def device(self) -> str: - return self._memory.device - - @property - def memory_type(self) -> str: - return self._memory.memory_type + work = Work(self, handle, local, remote) + self._works.add(work) + _live_endpoints.add(self) + # Retain the handle BEFORE posting. A post exception does not prove + # that nothing was submitted; endpoint.close() can retry cleanup. + work._state = self._backend.post(handle) + return work - @property - def closed(self) -> bool: - return self._closed + def notify(self, peer: str, message: bytes) -> None: + """Standalone notification: local acceptance, NOT remote acknowledgement.""" + with self._lock: + self._check_open() + if peer not in self._peers: + raise ValueError("connect the peer first") + self._backend.notify(peer, message) - def region( - self, - offset: int = 0, - nbytes: int | None = None, - *, - stride: int | None = None, - count: int | None = None, - ) -> Region: - self._endpoint._validate_registration(self) - offset, nbytes, stride, count = _validate_region_bounds( - self.nbytes, - offset, - nbytes, - stride, - count, - ) - if stride is not None and not self._endpoint.capabilities.strided_regions: - raise UnsupportedError( - f"backend {self._endpoint.backend!r} does not support strided regions" - ) - return Region( - self, - None, - offset=offset, - nbytes=nbytes, - device=self.device, - memory_type=self.memory_type, - remote=False, - stride=stride, - count=count, - _token=_CONSTRUCTION_TOKEN, - ) + def notifications(self) -> dict[str, list[bytes]]: + with self._lock: + self._check_open() + return self._backend.notifications() def close(self) -> None: - self._endpoint._close_registration(self) - - def _check_open(self) -> None: - if self._closed: - raise ClosedError(f"registration {self.name!r} is closed") - if self._creation_failed: - raise BackendFailureError( - f"registration {self.name!r} is retained for cleanup only" - ) - self._endpoint._check_open() - - def __enter__(self) -> Self: - self._check_open() - return self - - def __exit__(self, *args: object) -> None: - self.close() - - -class Peer: - """An imported remote endpoint and its named registrations.""" - - __slots__ = ( - "_active_sync_notification_sends", - "_closed", - "_creation_failed", - "_endpoint", - "_endpoint_id", - "_handle", - "_incarnation", - "_name", - "_registrations", - "_sync_notification_reservation_pool", - ) - - def __init__( - self, - endpoint: Endpoint, - metadata: EndpointMetadata, - handle: object, - *, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("Peer objects are created by Endpoint.import_peer") - self._endpoint = endpoint - self._endpoint_id = metadata.endpoint_id - self._name = metadata.name - self._incarnation = metadata.incarnation - self._handle = handle - self._creation_failed = False - self._registrations = { - registration.name: registration for registration in metadata.registrations - } - # MULTIPLE-mode synchronous sends have no Work object to retain this - # peer. The send frame strongly owns a pooled token while this set keeps - # only its weak identity. If asynchronous interruption prevents the - # cleanup epilogue, the dead weak reference cannot wedge peer closure. - self._active_sync_notification_sends: set[ - weakref.ReferenceType[_SyncNotificationReservation] - ] = set() - self._sync_notification_reservation_pool: list[_SyncNotificationReservation] = ( - [] - ) - self._closed = False - - @property - def endpoint_id(self) -> str: - return self._endpoint_id - - @property - def name(self) -> str: - return self._name + """Close in dependency order; a failure retains resources for retry. - @property - def incarnation(self) -> str: - return self._incarnation - - @property - def identity(self) -> EndpointIdentity: - return EndpointIdentity(self.endpoint_id, self.incarnation) - - @property - def closed(self) -> bool: - return self._closed - - def region( - self, - name: str, - offset: int = 0, - nbytes: int | None = None, - *, - stride: int | None = None, - count: int | None = None, - ) -> Region: - self._check_open() - try: - metadata = self._registrations[name] - except KeyError as error: - raise InvalidRegionError( - f"peer {self.name!r} has no registration {name!r}" - ) from error - offset, nbytes, stride, count = _validate_region_bounds( - metadata.nbytes, - offset, - nbytes, - stride, - count, - ) - if stride is not None and not self._endpoint.capabilities.strided_regions: - raise UnsupportedError( - f"backend {self._endpoint.backend!r} does not support strided regions" - ) - return Region( - self, - metadata, - offset=offset, - nbytes=nbytes, - device=metadata.device, - memory_type=metadata.memory_type, - remote=True, - stride=stride, - count=count, - _token=_CONSTRUCTION_TOKEN, - ) - - def unsafe_region( - self, - address: int, - nbytes: int, - *, - device: str, - memory_type: str, - stride: int | None = None, - count: int | None = None, - ) -> Region: - """Create a provider-validated remote raw-address region. - - This escape hatch exists for legacy protocols that exchange addresses. - The backend must validate that the span is covered by imported metadata. + The application must FIRST stop remote access to exported memory. + This method is not a distributed revoke, drain or shutdown protocol. """ + with self._lock: + for work in list(self._works): + work.close() + for catalog in list(self._catalogs): + catalog.close() + for peer in list(self._peers): + self.disconnect(peer) + for reg in list(self._registrations): + reg.close() + self._closed = True - self._check_open() - if not self._endpoint.capabilities.raw_spans: - raise UnsupportedError( - f"backend {self._endpoint.backend!r} does not support raw spans" - ) - _validate_positive_int(address, "address") - _validate_positive_int(nbytes, "nbytes") - _validate_nonempty_str(device, "device", InvalidRegionError) - _validate_nonempty_str(memory_type, "memory_type", InvalidRegionError) - stride, count = _validate_stride(stride, count) - if stride is not None and stride < nbytes: - raise InvalidRegionError("stride must be at least nbytes") - descriptor = _RemoteRawDescriptor( - address=address, - nbytes=_region_extent(nbytes, stride, count), - device=device, - memory_type=memory_type, - ) - return Region( - self, - descriptor, - offset=0, - nbytes=nbytes, - device=device, - memory_type=memory_type, - remote=True, - stride=stride, - count=count, - _token=_CONSTRUCTION_TOKEN, - ) - - def close(self) -> None: - self._endpoint._close_peer(self) - - def _check_open(self) -> None: - if self._closed: - raise ClosedError(f"peer {self.name!r} is closed") - if self._creation_failed: - raise BackendFailureError( - f"peer {self.name!r} is retained for cleanup only" - ) - self._endpoint._check_open() - def __enter__(self) -> Self: - self._check_open() - return self +class Registration: + def __init__(self, endpoint: Endpoint, handle: Any, owner: Any): + self._endpoint, self._handle, self._owner = endpoint, handle, owner - def __exit__(self, *args: object) -> None: - self.close() + def close(self) -> None: + ep = self._endpoint + with ep._lock: + if self not in ep._registrations: + return + # Deliberately coarse for the MVP: no region-to-catalog dependency graph. + if ep._catalogs: + raise RuntimeError("close catalogs before deregistering memory") + ep._backend.deregister(self._handle) + ep._registrations.remove(self) + self._owner = None -class TransferPlan: - """Reusable provider descriptors for one peer and one region layout. +class Catalog: + def __init__(self, endpoint: Endpoint, handle: Any, peer: str | None): + self._endpoint, self._handle, self.peer = endpoint, handle, peer - Indexed plans expose flattened region-major, logical-block catalogs. Each - Region contributes ``count`` blocks, or one block when ``count`` is absent. - """ + def close(self) -> None: + ep = self._endpoint + with ep._lock: + if self not in ep._catalogs: + return + if any(self in (w._local, w._remote) for w in ep._works): + raise RuntimeError("close dependent work before its catalog") + ep._backend.release_catalog(self._handle) + ep._catalogs.remove(self) - __slots__ = ( - "_bindings", - "_closed", - "_creation_failed", - "_endpoint", - "_handle", - "_indexed", - "_local", - "_local_block_count", - "_local_full_nonoverlap", - "_local_starts", - "_local_total_nbytes", - "_local_uniform_device", - "_local_uniform_nbytes", - "_origins", - "_remote", - "_remote_block_count", - "_remote_full_nonoverlap", - "_remote_starts", - "_remote_total_nbytes", - "_remote_uniform_nbytes", - "_reuse_tokens", - "_submitting", - "_works", - ) +class Work: def __init__( - self, - endpoint: Endpoint, - handle: object, - local: tuple[Region, ...], - remote: tuple[Region, ...], - indexed: bool, - *, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("TransferPlan objects are created by Endpoint.prepare") - self._endpoint = endpoint - self._handle = handle - self._local = local - self._remote = remote - self._indexed = indexed - self._local_block_count = _logical_block_count(local) - self._remote_block_count = _logical_block_count(remote) - self._local_starts = _catalog_starts(local) - self._remote_starts = _catalog_starts(remote) - self._local_total_nbytes = sum(region.total_nbytes for region in local) - self._remote_total_nbytes = sum(region.total_nbytes for region in remote) - self._local_uniform_nbytes = _uniform_region_nbytes(local) - self._remote_uniform_nbytes = _uniform_region_nbytes(remote) - self._local_uniform_device = _uniform_region_device(local) - self._local_full_nonoverlap = _regions_are_nonoverlapping(local) - self._remote_full_nonoverlap = _regions_are_nonoverlapping(remote) - self._origins = frozenset(region._origin for region in (*local, *remote)) - self._bindings: set[BoundTransfer] = set() - self._reuse_tokens = { - TransferOp.READ: _ReuseToken(), - TransferOp.WRITE: _ReuseToken(), - } - self._submitting: set[object] = set() - self._works: set[Work] = set() - self._creation_failed = False - self._closed = False - - @property - def indexed(self) -> bool: - return self._indexed - - @property - def local_block_count(self) -> int: - """Number of logical blocks in the local indexed catalog.""" - - return self._local_block_count - - @property - def remote_block_count(self) -> int: - """Number of logical blocks in the remote indexed catalog.""" - - return self._remote_block_count + self, endpoint: Endpoint, handle: Any, local: Catalog, remote: Catalog + ): + self._endpoint, self._handle = endpoint, handle + self._local, self._remote = local, remote + self._state, self._closed = "pending", False @property - def closed(self) -> bool: - return self._closed - - def bind( - self, - op: TransferOp | str, - local_indices: Sequence[int] | None = None, - remote_indices: Sequence[int] | None = None, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - ) -> BoundTransfer: - """Validate and retain one reusable submission shape. - - A binding amortizes operation, selection, notification, and CUDA-token - validation across repeated submissions. Each :meth:`BoundTransfer.submit` - still creates an independently owned :class:`Work`. - """ - - return self._endpoint._bind_transfer( - self, - op, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - ) - - def select_indices( - self, - op: TransferOp | str, - *, - local_indices: Sequence[int], - remote_indices: Sequence[int], - ) -> IndexSelection: - """Snapshot and validate one reusable dynamic indexed selection. - - C-contiguous native-int32 buffers use the packed fast input. Mutable - inputs are copied into one private Core-owned buffer, so later caller - mutation cannot change this selection. General sequences remain - supported through ``operator.index`` normalization. - """ - - return self._endpoint._select_plan_indices( - self, - op, - local_indices=local_indices, - remote_indices=remote_indices, - ) - - def submit_indices( - self, - op: TransferOp | str, - *, - local_indices: Sequence[int], - remote_indices: Sequence[int], - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - ) -> Work: - """Snapshot, validate, and submit one dynamic indexed transfer. - - This is the allocation-minimal one-shot form. It creates only Core's - private authenticated backend selection and retains that snapshot in - the returned Work. Use :meth:`select_indices` when the same validated - selection will be bound or reused. - """ - - return self._endpoint._submit_indexed_selection( - plan=self, - op=op, - local_indices=local_indices, - remote_indices=remote_indices, - selection=None, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - - def submit_indices_and_test( - self, - op: TransferOp | str, - *, - local_indices: Sequence[int], - remote_indices: Sequence[int], - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns: int | None = None, - ) -> tuple[Work, bool]: - """Compatibility wrapper returning whether the first state is terminal. - - The Boolean means terminal, including failure or cancellation. A zero - poll budget is post-only. Prefer :meth:`submit_indices_and_poll` when - the caller also needs the exact state; the wrapper does not issue an - additional provider observation. - """ - - work, state = self.submit_indices_and_poll( - op, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - return work, state in _TERMINAL_STATES - - def submit_indices_and_poll( - self, - op: TransferOp | str, - *, - local_indices: Sequence[int], - remote_indices: Sequence[int], - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns: int | None = None, - ) -> tuple[Work, WorkState]: - """Submit one dynamic selection and return its exact bounded-poll state. + def state(self) -> str: + """Last observed state without an extra provider query.""" + return self._state - A zero poll budget is post-only. Fused providers return the state from - the post boundary; other providers conservatively return ``PENDING`` - until at least one poll is requested. The returned Work remains the - authoritative owner in every state. - """ + def poll(self) -> str: + """Return pending/done/error. A polling exception is NOT completion.""" + ep = self._endpoint + with ep._lock: + if not self._closed and self._state == "pending": + self._state = ep._backend.poll(self._handle) + return self._state - _validate_bounded_poll_arguments( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=None, - allow_zero=True, - ) - fused = self._endpoint.capabilities.fused_prevalidated_submit_poll - work = self._endpoint._submit_indexed_selection( - plan=self, - op=op, - local_indices=local_indices, - remote_indices=remote_indices, - selection=None, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - initial_max_polls=max_polls if fused else None, - initial_timeout_ns=timeout_ns if fused else None, - ) - if fused: - return work, work._state - if max_polls == 0: - return work, WorkState.PENDING - return work, work.poll_bounded(max_polls=max_polls, timeout_ns=timeout_ns) + def wait(self, timeout: float | None = None) -> None: + deadline = None if timeout is None else time.monotonic() + timeout + while self.poll() == "pending": + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError("work is still live; timeout does not cancel it") + time.sleep(0) + if self._state != "done": + raise RuntimeError("transfer failed or was cancelled") def close(self) -> None: - self._endpoint._close_plan(self) - - def _check_open(self) -> None: - if self._closed: - raise ClosedError("transfer plan is closed") - if self._creation_failed: - raise BackendFailureError("transfer plan is retained for cleanup only") - self._endpoint._check_open() - - def __enter__(self) -> Self: - self._check_open() - return self - - def __exit__(self, *args: object) -> None: - self.close() - - -class IndexSelection: - """Immutable, exact-plan indexed selection with amortized validation. - - A selection is not an independently closeable resource: its exact plan may - be closed, which invalidates later ``submit`` and ``bind`` calls. Open Work - and BoundTransfer objects retain the selection (and therefore its packed - index snapshot) for as long as a provider may access it. - """ - - __slots__ = ( - "_backend_selection", - "_endpoint", - "_local_devices", - "_op", - "_plan", - "_reuse_token", - "_sealed", - ) - - def __init__( - self, - endpoint: Endpoint, - plan: TransferPlan, - op: TransferOp, - backend_selection: BackendIndexSelection, - local_devices: tuple[str, ...], - *, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("IndexSelection objects are created by TransferPlan") - set_attribute = object.__setattr__ - set_attribute(self, "_endpoint", endpoint) - set_attribute(self, "_plan", plan) - set_attribute(self, "_op", op) - set_attribute(self, "_backend_selection", backend_selection) - set_attribute(self, "_local_devices", local_devices) - # Allocate the provider-neutral plan/op/index-geometry token only if a - # selection-backed BoundTransfer is actually submitted. A one-shot - # selection and a binding used only to create a persistent slot do not - # need a request-pool identity. - set_attribute(self, "_reuse_token", None) - set_attribute(self, "_sealed", True) - - def __setattr__(self, name: str, value: object) -> None: - raise AttributeError("IndexSelection is immutable") - - def __delattr__(self, name: str) -> None: - raise AttributeError("IndexSelection is immutable") - - @property - def op(self) -> TransferOp: - return self._op - - @property - def local_count(self) -> int: - return len(self._backend_selection.local_indices) - - @property - def remote_count(self) -> int: - return len(self._backend_selection.remote_indices) - - @property - def selected_nbytes(self) -> int: - return self._backend_selection.selected_nbytes - - def submit( - self, - *, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - ) -> Work: - """Submit this selection with request-local completion arguments.""" - - return self._endpoint._submit_indexed_selection( - plan=None, - op=None, - local_indices=None, - remote_indices=None, - selection=self, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - - def submit_and_test( - self, - *, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns: int | None = None, - ) -> tuple[Work, bool]: - """Compatibility wrapper returning whether the first state is terminal. - - The Boolean has the same meaning as :meth:`Work.test`: ``True`` means - terminal, including failure or cancellation, rather than successful. - Prefer :meth:`submit_and_poll` when the exact state is needed. - """ - - work, state = self.submit_and_poll( - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - return work, state in _TERMINAL_STATES - - def submit_and_poll( - self, - *, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns: int | None = None, - ) -> tuple[Work, WorkState]: - """Submit and return the exact state from a bounded initial poll. - - A zero poll budget is post-only. On a provider without fused support, - Core conservatively returns ``PENDING`` because no post-derived state - has been observed. The returned Work remains authoritative. - """ - - _validate_bounded_poll_arguments( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=None, - allow_zero=True, - ) - if self._endpoint.capabilities.fused_prevalidated_submit_poll: - work = self._endpoint._submit_indexed_selection( - plan=None, - op=None, - local_indices=None, - remote_indices=None, - selection=self, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - initial_max_polls=max_polls, - initial_timeout_ns=timeout_ns, - ) - return work, work._state - - work = self.submit( - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - if max_polls == 0: - return work, WorkState.PENDING - return work, work.poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - - def bind( - self, - *, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - ) -> BoundTransfer: - """Bind completion arguments without revalidating this selection.""" - - return self._endpoint._bind_index_selection( - self, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - ) - - def _check_open(self) -> TransferPlan: - if self._plan._closed: - raise ClosedError("transfer plan is closed") - self._plan._check_open() - return self._plan - - -class BoundTransfer: - """A prevalidated, reusable transfer submission. - - Bindings retain their plan and canonical submission arguments. They are - Core-only resources: providers observe only the binding's opaque, stable - reuse token on submission. - """ - - __slots__ = ( - "_closed", - "_cuda_device", - "_cuda_ordering", - "_endpoint", - "_local_indices", - "_notification", - "_op", - "_plan", - "_ready_event", - "_remote_indices", - "_reuse_token", - "_selection", - "_slots", - "_stream", - "_submitting", - "_works", - ) - - def __init__( - self, - endpoint: Endpoint, - plan: TransferPlan, - op: TransferOp, - *, - local_indices: tuple[int, ...] | None, - remote_indices: tuple[int, ...] | None, - notification: bytes | None, - cuda_device: torch.device | None, - cuda_ordering: CudaOrderingMode, - stream: object | None, - ready_event: object | None, - selection: IndexSelection | None = None, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("BoundTransfer objects are created by TransferPlan.bind") - self._endpoint = endpoint - self._plan: TransferPlan | None = plan - self._op = op - self._local_indices = local_indices - self._remote_indices = remote_indices - self._notification = notification - self._cuda_device = cuda_device - self._cuda_ordering = cuda_ordering - self._stream = stream - self._ready_event = ready_event - self._selection = selection - # Selection-backed bindings share the selection's lazily allocated - # geometry token. Do not allocate a redundant binding token for them. - self._reuse_token: object | None = _ReuseToken() if selection is None else None - self._submitting: set[Work] = set() - self._works: set[Work] = set() - self._slots: set[TransferSlot] = set() - self._closed = False - - @property - def op(self) -> TransferOp: - return self._op - - @property - def closed(self) -> bool: - return self._closed - - def submit( - self, - *, - notification: bytes | bytearray | memoryview | None = cast( - Any, _UNSET_NOTIFICATION - ), - recovery_token: object | None = None, - ) -> Work: - """Submit the binding. - - Selection-backed bindings accept a request-local notification override - so one native indexed lowering can serve successive framework tags. - Ordinary bindings retain their original fully immutable contract. - """ - - return self._endpoint._submit_bound( - self, - notification=notification, - recovery_token=recovery_token, - ) - - def create_slot( - self, *, request_local_notifications: bool = False - ) -> TransferSlot | RequestLocalNotificationSlot: - """Create a persistent, single-flight request for this binding. - - The default returns the fixed-notification :class:`TransferSlot` and - preserves its minimal execute path. ``request_local_notifications=True`` - returns a :class:`RequestLocalNotificationSlot`, whose start operation - requires an explicit payload (or ``None`` to clear it) for every - generation. This opt-in mode is capability gated and requires a - binding with no fixed notification. - """ - - if not isinstance(request_local_notifications, bool): - raise TypeError("request_local_notifications must be a bool") - return self._endpoint._create_transfer_slot( - self, - request_local_notifications=request_local_notifications, - ) - - def close(self) -> None: - self._endpoint._close_bound_transfer(self) - - def _check_open(self) -> TransferPlan: - if self._closed: - raise ClosedError("bound transfer is closed") - plan = self._plan - if plan is None: - raise ClosedError("bound transfer is closed") - plan._check_open() - return plan - - def __enter__(self) -> Self: - self._check_open() - return self - - def __exit__(self, *args: object) -> None: - self.close() - - -class TransferSlot: - """Stable, mutable single-flight request for one ``BoundTransfer``. - - A slot owns at most one unrecycled generation. ``state`` is ``None`` while - idle and a :class:`WorkState` after ``start`` until terminal ``recycle``. - ``active`` deliberately remains true for a terminal generation: recycling - is the explicit boundary that permits the next start and prevents ABA-style - observations across generations. ``execute`` is the opt-in full-cycle form; - it publishes a clean successful generation directly back to idle. - """ - - __slots__ = ( - "_active", - "_backend_cancel", - "_backend_close", - "_backend_poll_bounded", - "_backend_poll_state", - "_backend_recycle", - "_backend_deferred_successful_recycle", - "_backend_execution_dispatch", - "_backend_generation_active", - "_caller_ready", - "_ready_event_synchronize", - "_fast_execute", - "_backend_slot", - "_backend_start", - "_backend_start_and_poll", - "_bound", - "_closed", - "_creation_failed", - "_endpoint", - "_error", - "_generation", - "_lock", - "_owner_thread", - "_start_error", - "_state", - ) - - def __init__( - self, - endpoint: Endpoint, - bound: BoundTransfer, - *, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("TransferSlot objects are created by BoundTransfer") - self._endpoint = endpoint - self._bound: BoundTransfer | None = bound - self._backend_slot: BackendRequestSlot | None = None - self._backend_execution_dispatch: object | None = None - self._backend_start: Callable[[], None] | None = None - self._backend_start_and_poll: Callable[..., WorkState | None] | None = None - self._backend_poll_bounded: Callable[..., WorkState | None] | None = None - self._backend_poll_state: Callable[[], WorkState | None] | None = None - self._backend_cancel: Callable[[], bool] | None = None - self._backend_recycle: Callable[[], None] | None = None - self._backend_close: Callable[[], None] | None = None - self._backend_deferred_successful_recycle = ( - endpoint.capabilities.deferred_successful_request_slot_recycle - ) - self._backend_generation_active = False - self._caller_ready = bound._cuda_ordering is CudaOrderingMode.CALLER_READY - # A BoundTransfer has already authenticated the token type and device. - # Pin one complete readiness operation per persistent slot. An explicit - # event contributes its bound synchronize method; a stream contributes - # one reusable, lazily materialized CUDA event plus cached capture, - # record, and synchronize callables. Every generation still performs - # capture detection, records after prior producer work, and host-waits, - # but it allocates neither a Python nor native CUDA event in steady state. - if self._caller_ready: - self._ready_event_synchronize = None - elif bound._stream is not None: - self._ready_event_synchronize = _make_reusable_cuda_stream_readiness( - bound._stream, - bound._cuda_device, - ) - else: - self._ready_event_synchronize = cast( - Callable[[], None], bound._ready_event.synchronize - ) - self._fast_execute: Callable[..., WorkState | None] | None = None - # SINGLE and CALLER_SERIALIZED use a genuinely lock-free path. ``None`` - # avoids even the Python context-manager calls paid by Core's general - # no-op lock. The latter contract requires whole-graph exclusion by its - # caller; it is not an independently thread-safe slot. - self._lock = ( - None - if endpoint._thread_mode - in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) - else threading.RLock() - ) - self._owner_thread = ( - endpoint._owner_thread - if endpoint._thread_mode is ThreadMode.SINGLE - else None - ) - self._generation = 0 - self._active = False - self._state: WorkState | None = None - self._error: BaseException | None = None - self._start_error: BaseException | None = None - self._creation_failed = False - self._closed = False - - @property - def generation(self) -> int: - """Number of starts owned by this slot, including retained failures.""" - - return self._generation - - @property - def active(self) -> bool: - """Whether this slot owns a generation that has not been recycled.""" - - return self._active - - @property - def closed(self) -> bool: - return self._closed - - @property - def state(self) -> WorkState | None: - return self._refresh() - - @property - def error(self) -> BaseException | None: - lock = self._lock - if lock is None: - self._check_access() - state = self._refresh_unlocked() - return None if state is None else self._error - with lock: - self._check_access() - state = self._refresh_unlocked() - return None if state is None else self._error - - @property - def completion_event(self) -> object | None: - lock = self._lock - if lock is None: - return self._completion_event_unlocked() - with lock: - return self._completion_event_unlocked() - - def start(self) -> None: - """Start one generation after re-establishing CUDA readiness. - - Once Core enters the provider ``start`` call, the generation is owned - even if an exception crosses that boundary. Such an exception is - retained and becomes visible as terminal failure; the slot cannot be - reused or closed until its provider state is terminal and recycled. - """ - - lock = self._lock - if lock is None: - self._start_unlocked() - return - with lock: - self._start_unlocked() - - def start_and_test(self) -> bool: - """Start and logically recycle a post-reported immediate success. - - ``True`` means the generation completed successfully and the public - slot is already idle/repostable. ``False`` means it remains active and - can be observed with :meth:`test_and_recycle`. Terminal failure is - raised and deliberately remains active for explicit error recovery. - """ - - fast_start_and_test = self._fast_execute - if fast_start_and_test is not None: - next_generation = self._generation + 1 - try: - if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError( - "transfer slot is retained for cleanup only" - ) - if ( - self._owner_thread is not None - and _GET_IDENT() != self._owner_thread - ): - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - if self._active: - raise BusyError( - "transfer slot already owns an unrecycled generation" - ) - - if not self._caller_ready: - ready_event_synchronize = self._ready_event_synchronize - if ready_event_synchronize is not None: - try: - ready_event_synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to wait for CUDA readiness" - ) from error - else: - bound = self._bound - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - try: - self._generation = next_generation - self._active = True - except BaseException: - self._active = False - self._generation = next_generation - 1 - self._state = None - self._error = None - self._start_error = None - self._backend_generation_active = False - raise - - state = fast_start_and_test(max_polls=0, timeout_ns=None) - except BaseException as error: - if self._generation != next_generation: - raise - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - if state is _WORK_COMPLETED: - self._active = False - return True - if state is _WORK_RUNNING: - # The fused provider has acknowledged this generation and - # reported the dominant QD-many nonterminal state. Avoid - # reclassifying the same enum through the generic observer; - # these two stores are exactly its RUNNING-state commit. - self._backend_generation_active = True - self._state = state - return False - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot start returned without " - "acknowledging a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - elif not isinstance(state, WorkState): - self._start_error = BackendFailureError( - "backend returned invalid fused request-slot state " - f"{state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - state = self._observe_backend_state_unlocked(state) - - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - start_and_poll = self._backend_start_and_poll - lock = self._lock - if start_and_poll is None: - if lock is None: - return self._start_and_test_fallback_unlocked() - with lock: - return self._start_and_test_fallback_unlocked() - if lock is None: - state = self._start_and_poll_unlocked(start_and_poll, None, max_polls=0) - expected_generation = self._generation - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - with lock: - state = self._start_and_poll_unlocked(start_and_poll, None, max_polls=0) - expected_generation = self._generation - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def _start_and_test_fallback_unlocked(self) -> bool: - """Keep fallback start/observe/finalize in one slot transaction.""" - - self._start_unlocked() - expected_generation = self._generation - state = self._refresh_unlocked() - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def test_and_recycle(self) -> bool: - """Observe once and logically recycle only successful completion. - - Returns ``False`` while the generation remains pending/running. - Successful completion returns ``True`` with the public slot idle. - Failure or cancellation is raised and remains active for explicit - recovery. Deferred-success providers retain the clean native receipt - for next-start/close without a provider recycle call. - """ - - lock = self._lock - if lock is None and self._backend_deferred_successful_recycle: - # Exact lock-free/deferred continuation path. Slot construction has - # already authenticated the lifetime graph and cached the provider - # dispatch. Keep the common PENDING/RUNNING/COMPLETED states on one - # straight-line poll without the generic refresh/observe/finalize - # frames; failure and malformed-provider states retain the full - # diagnostic path below. + """Release completed work; active cancellation is outside this MVP.""" + ep = self._endpoint + with ep._lock: if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError("transfer slot is retained for cleanup only") - if self._owner_thread is not None and _GET_IDENT() != self._owner_thread: - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - if not self._active: - raise BusyError("transfer slot has no active generation") - if self._state is WorkState.COMPLETED: - # A preceding state/test observation may already have cached - # successful completion. Commit Core-idle directly; polling - # the provider again cannot add information and penalizes the - # normal test-then-recycle composition. - self._fast_execute = None - self._backend_generation_active = False - self._error = None - self._start_error = None - self._state = None - # Deferred success leaves a durable provider receipt, so all - # private Core state can be clean before the public idle commit. - # This prevents an asynchronous exception from exposing stale - # terminal state beside a subclass-specific fast dispatcher. - self._active = False - self._restore_fast_execute_unlocked() - return True - if self._state in _TERMINAL_STATES: - self._raise_for_state() - poll_state = self._backend_poll_state - if poll_state is None: - raise BackendFailureError( - "transfer slot has no backend state-poll operation" - ) - try: - state = poll_state() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._endpoint.backend!r} failed to query " - "persistent request state" - ) from error - if state is WorkState.COMPLETED and self._start_error is None: - self._state = WorkState.COMPLETED - self._backend_generation_active = True - self._fast_execute = None - self._backend_generation_active = False - self._error = None - self._start_error = None - self._state = None - # Keep the deferred-success idle commit last for the same - # interruption-safety reason as the cached-completion branch. - self._active = False - self._restore_fast_execute_unlocked() - return True - if state is None: - if self._start_error is None: - self._start_error = BackendFailureError( - "backend request-slot start returned without " - "acknowledging a generation" - ) - self._error = self._start_error - self._state = WorkState.FAILED - self._raise_for_state() - if not isinstance(state, WorkState): - raise BackendFailureError( - f"backend returned invalid transfer slot state {state!r}" - ) - if self._state is WorkState.RUNNING and state is WorkState.PENDING: - raise BackendFailureError( - "backend transfer slot state regressed from running to pending" - ) - self._backend_generation_active = True - if state is WorkState.PENDING or state is WorkState.RUNNING: - self._state = state - return False - self._observe_backend_state_unlocked(state) - self._raise_for_state() - - if lock is None: - self._check_access() - if not self._active: - raise BusyError("transfer slot has no active generation") - expected_generation = self._generation - state = self._refresh_unlocked() - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - with lock: - self._check_access() - if not self._active: - raise BusyError("transfer slot has no active generation") - expected_generation = self._generation - state = self._refresh_unlocked() - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def test(self) -> bool: - state = self._refresh() - return state in _TERMINAL_STATES if state is not None else False - - def wait(self, timeout: float | None = None) -> TransferSlot: - deadline_ns = _transfer_slot_deadline_ns(timeout) - lock = self._lock - if lock is None: - expected_generation = self._generation_for_wait_unlocked() - else: - with lock: - expected_generation = self._generation_for_wait_unlocked() - return self._wait_until( - deadline_ns, - spin=False, - expected_generation=expected_generation, - ) - - def start_and_wait(self, timeout: float | None = None) -> TransferSlot: - """Start one generation and wait, using bounded provider fusion. - - The optional timeout is a total host-side budget beginning at method - entry. It cannot preempt a blocking CUDA-readiness or provider call. - Timeout leaves the generation active so it can be observed, cancelled, - and explicitly recycled after reaching terminal state. - """ - - deadline_ns = _transfer_slot_deadline_ns(timeout) - start_and_poll = self._backend_start_and_poll - poll_bounded = self._backend_poll_bounded - if start_and_poll is None: - lock = self._lock - if lock is None: - self._start_unlocked() - expected_generation = self._generation - else: - with lock: - self._start_unlocked() - expected_generation = self._generation - if poll_bounded is not None: - return self._wait_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_until( - deadline_ns, spin=True, expected_generation=expected_generation - ) - - lock = self._lock - if lock is None: - state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self - if state in _TERMINAL_STATES: - return self._raise_for_state() - else: - with lock: - state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self - if state in _TERMINAL_STATES: - return self._raise_for_state() - - # A provider is allowed to exhaust its small bounded poll budget while - # work remains active. Do not retain the per-slot lock across the - # potentially long ordinary wait. - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - # Preserve the zero-timeout no-extra-poll contract without letting - # a concurrent recycle/restart be misreported as a timeout. - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if poll_bounded is not None: - return self._wait_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_until( - deadline_ns, - spin=True, - expected_generation=expected_generation, - ) - - def execute(self, timeout: float | None = None) -> int: - """Start, wait, and logically recycle one successful generation. - - The returned integer identifies the completed generation. Success makes - the public slot idle before returning, so ``state`` is ``None`` and a - later generation may start immediately. Failure, cancellation, and - timeout deliberately retain the active generation for observation, - cancellation, and explicit recycle. - - Providers advertising deferred successful recycle may fold their - successful provider-side recycle into the next start or close boundary. - Other providers are recycled before Core publishes the slot as idle. - """ - - fast_execute = self._fast_execute - - # Specialize the dominant small-transfer lifecycle as one straight-line - # Core transaction: lock-free ownership, an unbounded fused provider - # call, and a provider that retains clean successful receipts. This - # preserves the generic transaction's checks and interruption - # boundaries while avoiding four Python helper frames on immediate - # completion. - if timeout is None and fast_execute is not None: - next_generation = self._generation + 1 - try: - if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError( - "transfer slot is retained for cleanup only" - ) - if ( - self._owner_thread is not None - and _GET_IDENT() != self._owner_thread - ): - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - if self._active: - raise BusyError( - "transfer slot already owns an unrecycled generation" - ) - - if not self._caller_ready: - ready_event_synchronize = self._ready_event_synchronize - if ready_event_synchronize is not None: - try: - ready_event_synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to wait for CUDA readiness" - ) from error - else: - bound = self._bound - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - try: - # Idle already guarantees clear diagnostics and no active - # provider generation. `_active` is the sole publication - # commit; the provider dispatch owns durable request - # state after entry, so this synchronous SINGLE path needs - # no transient Core PENDING store. - self._generation = next_generation - self._active = True - except BaseException: - self._active = False - self._generation = next_generation - 1 - self._state = None - self._error = None - self._start_error = None - self._backend_generation_active = False - raise - - state = fast_execute( - max_polls=_FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns=None, - ) - except BaseException as error: - if self._generation != next_generation: - raise - # Once Core publishes the generation, a provider call may have - # launched native work. Preserve the same ambiguous-start - # recovery used by the generic fused path. - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - if state is _WORK_COMPLETED: - # The provider completion receipt is already durable. Before - # idle commit, an active Core generation rediscovers it; - # after commit, Core is already clean-idle. No transient - # Core state stores are needed on this exact path. - self._active = False - return next_generation - if state is _WORK_RUNNING: - # The exact RUNNING singleton is already a valid provider - # acknowledgement. Publish provider ownership before the - # observable state, then enter the cached continuation - # directly instead of reclassifying the initial result. - self._backend_generation_active = True - self._state = state - poll_bounded = self._backend_poll_bounded - if poll_bounded is not None: - return self._wait_execute_bounded_until( - None, - expected_generation=next_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - None, expected_generation=next_generation - ) - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot start returned without " - "acknowledging a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - elif not isinstance(state, WorkState): - self._start_error = BackendFailureError( - "backend returned invalid fused request-slot state " - f"{state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - state = self._observe_backend_state_unlocked(state) - - expected_generation = self._generation - if state in _TERMINAL_STATES: - self._raise_for_state() - poll_bounded = self._backend_poll_bounded - if poll_bounded is not None: - return self._wait_execute_bounded_until( - None, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - None, expected_generation=expected_generation - ) - - # Non-fast paths preserve the shared finite-timeout validator. Keeping - # this below the fast return avoids a deadline local on the dominant - # unbounded deferred SINGLE completion. - deadline_ns = None if timeout is None else _transfer_slot_deadline_ns(timeout) - start_and_poll = self._backend_start_and_poll - lock = self._lock - poll_bounded = self._backend_poll_bounded - if start_and_poll is None: - if lock is None: - self._start_unlocked() - expected_generation = self._generation - else: - with lock: - self._start_unlocked() - expected_generation = self._generation - if poll_bounded is not None: - return self._wait_execute_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - deadline_ns, expected_generation=expected_generation - ) - - if lock is None: - state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) - if state is WorkState.COMPLETED: - return self._finish_execute_success_unlocked(self._generation) - expected_generation = self._generation - if state in _TERMINAL_STATES: - self._raise_for_state() - else: - with lock: - state = self._start_and_poll_unlocked(start_and_poll, deadline_ns) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self._finish_execute_success_unlocked(expected_generation) - if state in _TERMINAL_STATES: - self._raise_for_state() - - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if poll_bounded is not None: - return self._wait_execute_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - deadline_ns, expected_generation=expected_generation - ) - - def _wait_until( - self, - deadline_ns: int | None, - *, - spin: bool, - expected_generation: int, - ) -> TransferSlot: - polls = 0 - while True: - if self._observe_wait_generation(expected_generation): - return self - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - raise TransferTimeoutError("transfer slot timed out") - if self._endpoint._progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - polls += 1 - if not spin and polls >= 16: - time.sleep(0 if polls < 128 else 0.0001) - - def _wait_execute_until( - self, - deadline_ns: int | None, - *, - expected_generation: int, - ) -> int: - """Spin until success and publish Core-idle in the observation lock.""" - - while True: - lock = self._lock - if lock is None: - completed_generation = self._observe_execute_generation_unlocked( - expected_generation - ) - else: - with lock: - completed_generation = self._observe_execute_generation_unlocked( - expected_generation - ) - if completed_generation is not None: - return completed_generation - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if self._endpoint._progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - - def _wait_bounded_until( - self, - deadline_ns: int | None, - *, - expected_generation: int, - poll_bounded: Callable[..., WorkState | None], - ) -> TransferSlot: - lock = self._lock - while True: - if lock is None: - terminal = self._poll_bounded_generation_unlocked( - expected_generation, deadline_ns, poll_bounded - ) - else: - with lock: - terminal = self._poll_bounded_generation_unlocked( - expected_generation, deadline_ns, poll_bounded - ) - if terminal: - return self - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - raise TransferTimeoutError("transfer slot timed out") - if self._endpoint._progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - - def _wait_execute_bounded_until( - self, - deadline_ns: int | None, - *, - expected_generation: int, - poll_bounded: Callable[..., WorkState | None], - ) -> int: - """Run bounded polls and publish successful recycle under the slot lock.""" - - lock = self._lock - while True: - if lock is None: - terminal = self._poll_bounded_generation_unlocked( - expected_generation, deadline_ns, poll_bounded - ) - if terminal: - return self._finish_execute_success_unlocked(expected_generation) - else: - with lock: - terminal = self._poll_bounded_generation_unlocked( - expected_generation, deadline_ns, poll_bounded - ) - if terminal: - return self._finish_execute_success_unlocked( - expected_generation - ) - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if self._endpoint._progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - - def _poll_bounded_generation_unlocked( - self, - expected_generation: int, - deadline_ns: int | None, - poll_bounded: Callable[..., WorkState | None], - ) -> bool: - self._validate_wait_generation_unlocked(expected_generation) - if self._state is WorkState.COMPLETED: - return True - if self._state in _TERMINAL_STATES: - self._raise_for_state() - configured_timeout_ns = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS - if deadline_ns is None: - timeout_ns = configured_timeout_ns - else: - remaining_ns = max(0, deadline_ns - time.monotonic_ns()) - timeout_ns = ( - remaining_ns - if configured_timeout_ns is None - else min(configured_timeout_ns, remaining_ns) - ) - if self._lock is None: - try: - state = poll_bounded( - max_polls=_FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - timeout_ns=timeout_ns, - timeout_check_interval=( - _FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL - ), - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._endpoint.backend!r} failed during bounded " - "persistent request polling" - ) from error - else: - state = self._endpoint._backend_request_slot_poll_bounded( - poll_bounded, - max_polls=_FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - timeout_ns=timeout_ns, - timeout_check_interval=( - _FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL - ), - ) - state = self._observe_backend_state_unlocked(state) - if state is WorkState.COMPLETED: - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def cancel(self) -> bool: - lock = self._lock - if lock is None: - return self._cancel_unlocked() - with lock: - return self._cancel_unlocked() - - def recycle(self) -> None: - """Return a terminal generation to idle without changing its shape.""" - - lock = self._lock - if lock is None: - self._recycle_unlocked() - return - with lock: - self._recycle_unlocked() - - def close(self) -> None: - lock = self._lock - if lock is None: - self._endpoint._close_transfer_slot(self) - return - with lock: - self._endpoint._close_transfer_slot(self) - - def _adopt_backend_slot( - self, - backend_slot: BackendRequestSlot, - *, - execution_dispatch: object | None = None, - ) -> None: - if not isinstance(backend_slot, BackendRequestSlot): - raise BackendFailureError( - "backend must adopt a BackendRequestSlot instance" - ) - lock = self._lock - if lock is None: - self._adopt_backend_slot_unlocked( - backend_slot, execution_dispatch=execution_dispatch - ) - return - with lock: - self._adopt_backend_slot_unlocked( - backend_slot, execution_dispatch=execution_dispatch - ) - - def _refresh(self) -> WorkState | None: - lock = self._lock - if lock is None: - self._check_access() - return self._refresh_unlocked() - with lock: - self._check_access() - return self._refresh_unlocked() - - def _start_unlocked(self) -> None: - previous_generation = self._generation - try: - self._begin_generation_unlocked() - start = self._backend_start - if start is None: - self._start_error = BackendFailureError( - "transfer slot has no adopted backend request" - ) return - if self._lock is None: - result = start() - else: - result = self._endpoint._backend_request_slot_start(start) - # A normal return acknowledges that provider start was entered, - # even when the provider then violates the required None return. - self._backend_generation_active = True - if result is not None: - raise BackendFailureError("backend request-slot start must return None") - except BaseException as error: - if self._generation == previous_generation: - # CUDA readiness or an interrupted pre-commit publication did - # not create a generation and must remain a direct exception. - raise - # Native launch is ambiguous once the provider call has begun. - # Preserve ownership; a conforming provider exposes terminal state - # or permits cancellation rather than stranding this generation. - self._start_error = error - - def _begin_generation_unlocked(self) -> None: - bound = self._check_access() - if self._active: - raise BusyError("transfer slot already owns an unrecycled generation") - - if self._state is not None: - self._backend_generation_active = False - self._error = None - self._start_error = None - self._state = None - - # CUDA readiness failure is unambiguously before provider activity, so - # it does not consume a generation. This runs outside endpoint state - # and provider locks and is repeated for every start. - if not self._caller_ready: - ready_event_synchronize = self._ready_event_synchronize - if ready_event_synchronize is not None: - try: - ready_event_synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to wait for CUDA readiness" - ) from error - else: - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - previous_generation = self._generation - try: - # A generic/mixed generation temporarily disables the exact execute - # path. Its idle transition restores the cached callable only after - # all private state is clean. CUDA readiness precedes this store, so - # an unambiguous readiness failure preserves fast eligibility. - self._fast_execute = None - # ``_active`` is the publication commit and is deliberately last. - # An interrupt before it leaves an idle slot; an interrupt after it - # leaves a complete, self-consistent Core generation that the - # caller can observe and recycle without touching an idle provider. - self._state = WorkState.PENDING - self._error = None - self._start_error = None - self._backend_generation_active = False - self._generation = previous_generation + 1 - self._active = True - except BaseException: - self._active = False - self._generation = previous_generation - self._state = None - self._error = None - self._start_error = None - self._backend_generation_active = False - self._restore_fast_execute_unlocked() - raise - - def _start_and_poll_unlocked( - self, - start_and_poll: Callable[..., WorkState | None], - deadline_ns: int | None, - *, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - ) -> WorkState: - previous_generation = self._generation - try: - self._begin_generation_unlocked() - timeout_ns = ( - None - if deadline_ns is None - else min( - _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS, - max(0, deadline_ns - time.monotonic_ns()), - ) - ) - if self._lock is None: - state = start_and_poll( - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - else: - state = self._endpoint._backend_request_slot_start_and_poll( - start_and_poll, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - except BaseException as error: - if self._generation == previous_generation: - raise - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot start returned without acknowledging " - "a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - if not isinstance(state, WorkState): - self._start_error = BackendFailureError( - f"backend returned invalid fused request-slot state {state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - return self._observe_backend_state_unlocked(state) - - def _completion_event_unlocked(self) -> object | None: - self._check_access() - if not self._active: - raise BusyError("transfer slot has no active generation") - if not self._endpoint.capabilities.cuda_completion_event: - return None - backend_slot = self._require_backend_slot() - return self._endpoint._backend_request_slot_completion_event(backend_slot) - - def _generation_for_wait_unlocked(self) -> int: - self._check_access() - if not self._active: - raise BusyError("transfer slot has no active generation") - return self._generation - - def _observe_wait_generation(self, expected_generation: int) -> bool: - lock = self._lock - if lock is None: - return self._observe_wait_generation_unlocked(expected_generation) - with lock: - return self._observe_wait_generation_unlocked(expected_generation) - - def _validate_wait_generation(self, expected_generation: int) -> None: - lock = self._lock - if lock is None: - self._validate_wait_generation_unlocked(expected_generation) - return - with lock: - self._validate_wait_generation_unlocked(expected_generation) - - def _validate_wait_generation_unlocked(self, expected_generation: int) -> None: - self._check_access() - if not self._active or self._generation != expected_generation: - raise BusyError( - f"transfer slot generation {expected_generation} changed while " - "waiting" - ) - - def _observe_wait_generation_unlocked(self, expected_generation: int) -> bool: - self._check_access() - if not self._active or self._generation != expected_generation: - raise BusyError( - f"transfer slot generation {expected_generation} changed while " - "waiting" - ) - state = self._refresh_unlocked() - if state in _TERMINAL_STATES: - self._raise_for_state() - return True - return False - - def _observe_execute_generation_unlocked( - self, expected_generation: int - ) -> int | None: - self._validate_wait_generation_unlocked(expected_generation) - state = self._refresh_unlocked() - if state is WorkState.COMPLETED: - return self._finish_execute_success_unlocked(expected_generation) - if state in _TERMINAL_STATES: - self._raise_for_state() - return None - - def _cancel_unlocked(self) -> bool: - self._check_access() - if not self._active: - return False - if self._refresh_unlocked() in _TERMINAL_STATES: - return self._state is WorkState.CANCELLED - cancel = self._backend_cancel - if cancel is None: - raise BackendFailureError("transfer slot has no adopted backend request") - cancelled = self._endpoint._backend_request_slot_cancel(cancel) - if not isinstance(cancelled, bool): - raise BackendFailureError("backend slot cancellation must return a bool") - self._refresh_unlocked() - return cancelled and self._state is WorkState.CANCELLED - - def _recycle_unlocked(self) -> None: - self._check_access() - if not self._active: - raise BusyError("transfer slot has no generation to recycle") - self._fast_execute = None - if self._refresh_unlocked() not in _TERMINAL_STATES: - raise BusyError("cannot recycle a non-terminal transfer slot") - if self._backend_generation_active: - recycle = self._backend_recycle - if recycle is None: - raise BackendFailureError( - "transfer slot has no adopted backend request" - ) - if self._lock is None: - try: - result = recycle() - if result is not None: - raise BackendFailureError( - "backend request-slot recycle must return None" - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._endpoint.backend!r} failed to recycle a " - "persistent request" - ) from error - else: - self._endpoint._backend_request_slot_recycle(recycle) - # ``_active = False`` is the sole public idle-state commit. Private - # terminal diagnostics may remain briefly if an asynchronous exception - # lands after it; public state/error ignore them while idle, and the - # next generation overwrites them before its own active commit. - self._active = False - self._backend_generation_active = False - self._error = None - self._start_error = None - # Clear the dirty-idle marker only after every other private field. - self._state = None - self._restore_fast_execute_unlocked() - - def _finish_execute_success_unlocked(self, expected_generation: int) -> int: - """Commit a successful full-cycle generation while holding its lock.""" - - # Every caller has already validated access and generation while holding - # this slot's serialization boundary. Retain the cheap inline ABA guard - # without repeating `_check_access()` on every successful generation. - if not self._active or self._generation != expected_generation: - raise BusyError( - f"transfer slot generation {expected_generation} changed while " - "waiting" - ) - if self._state is not WorkState.COMPLETED: - raise BackendFailureError( - "transfer slot execute success requires completed provider state" - ) - self._fast_execute = None - if ( - self._backend_generation_active - and not self._backend_deferred_successful_recycle - ): - recycle = self._backend_recycle - if recycle is None: - raise BackendFailureError( - "transfer slot has no adopted backend request" - ) - if self._lock is None: - try: - result = recycle() - if result is not None: - raise BackendFailureError( - "backend request-slot recycle must return None" - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._endpoint.backend!r} failed to recycle a " - "persistent request" - ) from error - else: - self._endpoint._backend_request_slot_recycle(recycle) - - completed_generation = self._generation - # The provider retains an observable COMPLETED receipt when its recycle - # is deferred. ``_active = False`` is the sole public idle commit; the - # private terminal state is cleared afterward. Restore the exact fast - # callable only after every private field is clean. - self._active = False - self._backend_generation_active = False - self._error = None - self._start_error = None - self._state = None - self._restore_fast_execute_unlocked() - return completed_generation - - def _restore_fast_execute_unlocked(self) -> None: - """Publish exact-execute eligibility after a completely clean idle.""" - - start_and_poll = self._backend_start_and_poll - self._fast_execute = ( - start_and_poll - if ( - self._lock is None - and start_and_poll is not None - and self._backend_deferred_successful_recycle - ) - else None - ) - - def _adopt_backend_slot_unlocked( - self, - backend_slot: BackendRequestSlot, - *, - execution_dispatch: object | None, - ) -> None: - if self._closed: - raise BackendFailureError( - "backend attempted to adopt an expired transfer slot" - ) - if self._backend_slot is not None: - raise BackendFailureError( - "backend attempted to adopt transfer slot more than once" - ) - # Allocate every bound method before publishing `_backend_slot` as the - # adoption commit marker. An asynchronous interruption therefore leaves - # either a wholly unadopted slot or a fully usable one, never a public - # slot with only part of its steady-state dispatch table installed. - dispatch = backend_slot if execution_dispatch is None else execution_dispatch - start = dispatch.start - poll_state = dispatch.poll_state - start_and_poll = ( - dispatch.start_and_poll - if self._endpoint.capabilities.fused_request_slot_start_poll - else None - ) - poll_bounded = ( - dispatch.poll_bounded - if self._endpoint.capabilities.fused_request_slot_poll - else None - ) - cancel = dispatch.cancel - recycle = dispatch.recycle - close = backend_slot.close - operations = ( - ("start", start), - ("poll_state", poll_state), - ("cancel", cancel), - ("recycle", recycle), - ("close", close), - ) - if self._endpoint.capabilities.fused_request_slot_start_poll: - operations += (("start_and_poll", start_and_poll),) - if self._endpoint.capabilities.fused_request_slot_poll: - operations += (("poll_bounded", poll_bounded),) - for name, operation in operations: - if not callable(operation): - raise BackendFailureError( - f"backend request-slot {name} operation must be callable" - ) - self._backend_start = start - self._backend_poll_state = poll_state - self._backend_start_and_poll = start_and_poll - self._backend_poll_bounded = poll_bounded - self._backend_cancel = cancel - self._backend_recycle = recycle - self._backend_close = close - self._backend_execution_dispatch = dispatch - self._fast_execute = ( - start_and_poll - if ( - self._lock is None - and start_and_poll is not None - and self._backend_deferred_successful_recycle - ) - else None - ) - self._backend_slot = backend_slot - - def _refresh_unlocked(self) -> WorkState | None: - if not self._active: - return None - if self._state is WorkState.COMPLETED: - return self._state - if self._state in _TERMINAL_STATES: - return self._state - poll_state = self._backend_poll_state - if poll_state is None: - raise BackendFailureError( - "transfer slot has no backend state-poll operation" - ) - if self._lock is None: - try: - state = poll_state() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._endpoint.backend!r} failed to query " - "persistent request state" - ) from error - else: - state = self._endpoint._backend_request_slot_state(poll_state) - if state is None: - if self._start_error is None: - self._start_error = BackendFailureError( - "backend request-slot start returned without acknowledging " - "a generation" - ) - self._error = self._start_error - self._state = WorkState.FAILED - return self._state - return self._observe_backend_state_unlocked(state) - - def _observe_backend_state_unlocked(self, state: object) -> WorkState: - if not isinstance(state, WorkState): - raise BackendFailureError( - f"backend returned invalid transfer slot state {state!r}" - ) - if self._state is WorkState.RUNNING and state is WorkState.PENDING: - raise BackendFailureError( - "backend transfer slot state regressed from running to pending" - ) - if not self._backend_generation_active: - self._backend_generation_active = True - if state is WorkState.COMPLETED and self._start_error is None: - self._state = state - return state - observed_state = state - observed_error: BaseException | None = None - if state in _TERMINAL_STATES and self._start_error is not None: - observed_state = WorkState.FAILED - observed_error = self._start_error - elif state is WorkState.FAILED: - try: - backend_slot = self._require_backend_slot() - error = self._endpoint._backend_request_slot_error(backend_slot) - except Exception as observation_error: - observed_error = observation_error - else: - if error is None: - error = BackendFailureError( - "transfer backend reported slot failure without an error" - ) - elif not isinstance(error, BaseException): - error = BackendFailureError( - "transfer backend returned a non-exception slot error" - ) - observed_error = error - elif state is WorkState.CANCELLED: - observed_error = TransferCancelledError("transfer slot was cancelled") - - # Error construction/provider inspection happens first; terminal state - # is the commit marker. A BaseException during error discovery leaves - # the prior nonterminal state, so the next observation retries instead - # of permanently caching a terminal state with no diagnostic. - if observed_state in _TERMINAL_STATES: - self._error = observed_error - self._state = observed_state - return observed_state - - def _raise_for_state(self) -> TransferSlot: - if self._state is WorkState.COMPLETED: - return self - if self._error is None: - self._error = BackendFailureError( - f"transfer slot reached {self._state.value} without an error" - ) - if isinstance(self._error, TransferError): - raise self._error - raise BackendFailureError("transfer slot failed") from self._error - - def _require_backend_slot(self) -> BackendRequestSlot: - backend_slot = self._backend_slot - if backend_slot is None: - raise BackendFailureError("transfer slot has no adopted backend request") - return backend_slot - - def _check_access(self) -> BoundTransfer: - if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError("transfer slot is retained for cleanup only") - bound = self._bound - if bound is None: - raise ClosedError("transfer slot is closed") - owner_thread = self._owner_thread - if owner_thread is not None and threading.get_ident() != owner_thread: - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - # An open slot prevents its BoundTransfer, plan, and endpoint from - # closing. Creation validates that graph once, so polling need not walk - # it again on every request-state query. - return bound - - def __enter__(self) -> Self: - self._check_access() - return self - - def __exit__(self, *args: object) -> None: - self.close() - - -class RequestLocalNotificationSlot(TransferSlot): - """Persistent request with an explicit notification per generation. - - The immutable operation, descriptors, index selection, CUDA-ordering - contract, and native request storage are inherited from ``TransferSlot``. - Only the attached notification varies. Every start must state either a - non-empty bytes-like payload or ``None``; the latter is an explicit native - clear, not "reuse the previous value". The distinct type and methods keep - all validation and dispatch branches out of fixed ``TransferSlot.execute``. - """ - - __slots__ = ( - "_backend_start_with_notification", - "_backend_start_and_poll_with_notification", - "_notification_fast_execute", - ) - - def __init__( - self, - endpoint: Endpoint, - bound: BoundTransfer, - *, - _token: object, - ) -> None: - super().__init__(endpoint, bound, _token=_token) - self._backend_start_with_notification: Callable[[bytes | None], None] | None = ( - None - ) - self._backend_start_and_poll_with_notification: ( - Callable[..., WorkState | None] | None - ) = None - self._notification_fast_execute: Callable[..., WorkState | None] | None = None - - def start( - self, - *, - notification: bytes | bytearray | memoryview | None, - ) -> None: - """Start one generation with an explicit notification or clear.""" - - if notification is not None: - if type(notification) is bytes: - if not notification: - raise ValueError( - "an attached notification payload must be non-empty" - ) - if len(notification) > _MAX_NOTIFICATION_BYTES: - raise ValueError( - "notification exceeds the " - f"{_MAX_NOTIFICATION_BYTES}-byte size limit" - ) - else: - notification = _validate_request_slot_notification(notification) - lock = self._lock - if lock is None: - self._start_with_notification_unlocked(notification) - return - with lock: - self._start_with_notification_unlocked(notification) - - def start_and_wait( - self, - *, - notification: bytes | bytearray | memoryview | None, - timeout: float | None = None, - ) -> RequestLocalNotificationSlot: - """Start with an explicit notification and wait for terminal state.""" - - if notification is not None: - if type(notification) is bytes: - if not notification: - raise ValueError( - "an attached notification payload must be non-empty" - ) - if len(notification) > _MAX_NOTIFICATION_BYTES: - raise ValueError( - "notification exceeds the " - f"{_MAX_NOTIFICATION_BYTES}-byte size limit" - ) - else: - notification = _validate_request_slot_notification(notification) - deadline_ns = _transfer_slot_deadline_ns(timeout) - start_and_poll = self._backend_start_and_poll_with_notification - poll_bounded = self._backend_poll_bounded - if start_and_poll is None: - lock = self._lock - if lock is None: - self._start_with_notification_unlocked(notification) - expected_generation = self._generation - else: - with lock: - self._start_with_notification_unlocked(notification) - expected_generation = self._generation - if poll_bounded is not None: - self._wait_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - else: - self._wait_until( - deadline_ns, - spin=True, - expected_generation=expected_generation, - ) - return self - - lock = self._lock - if lock is None: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, notification, deadline_ns - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self - if state in _TERMINAL_STATES: - return cast(RequestLocalNotificationSlot, self._raise_for_state()) - else: - with lock: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, notification, deadline_ns - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self - if state in _TERMINAL_STATES: - return cast(RequestLocalNotificationSlot, self._raise_for_state()) - - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if poll_bounded is not None: - self._wait_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - else: - self._wait_until( - deadline_ns, - spin=True, - expected_generation=expected_generation, - ) - return self - - def start_and_test( - self, - *, - notification: bytes | bytearray | memoryview | None, - ) -> bool: - """Start with an explicit tag/clear and consume immediate success.""" - - if notification is not None: - if type(notification) is bytes: - if not notification: - raise ValueError( - "an attached notification payload must be non-empty" - ) - if len(notification) > _MAX_NOTIFICATION_BYTES: - raise ValueError( - "notification exceeds the " - f"{_MAX_NOTIFICATION_BYTES}-byte size limit" - ) - else: - notification = _validate_request_slot_notification(notification) - - fast_start_and_test = self._notification_fast_execute - if fast_start_and_test is not None: - next_generation = self._generation + 1 - try: - if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError( - "transfer slot is retained for cleanup only" - ) - if ( - self._owner_thread is not None - and _GET_IDENT() != self._owner_thread - ): - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - if self._active: - raise BusyError( - "transfer slot already owns an unrecycled generation" - ) - - if not self._caller_ready: - ready_event_synchronize = self._ready_event_synchronize - if ready_event_synchronize is not None: - try: - ready_event_synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to wait for CUDA readiness" - ) from error - else: - bound = self._bound - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - try: - self._generation = next_generation - self._active = True - except BaseException: - self._active = False - self._generation = next_generation - 1 - self._state = None - self._error = None - self._start_error = None - self._backend_generation_active = False - raise - - state = fast_start_and_test( - notification, - max_polls=0, - timeout_ns=None, - ) - except BaseException as error: - if self._generation != next_generation: - raise - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - if state is _WORK_COMPLETED: - self._active = False - return True - if state is _WORK_RUNNING: - # Match the fixed-notification slot's exact QD-many path: - # the provider already authenticated the enum and Core - # needs only to publish the acknowledged RUNNING state. - self._backend_generation_active = True - self._state = state - return False - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot notification start returned " - "without acknowledging a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - elif not isinstance(state, WorkState): - self._start_error = BackendFailureError( - "backend returned invalid fused request-slot state " - f"{state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - state = self._observe_backend_state_unlocked(state) - - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - start_and_poll = self._backend_start_and_poll_with_notification - lock = self._lock - if start_and_poll is None: - if lock is None: - return self._start_and_test_with_notification_fallback_unlocked( - notification - ) - with lock: - return self._start_and_test_with_notification_fallback_unlocked( - notification - ) - if lock is None: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, - notification, - None, - max_polls=0, - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - with lock: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, - notification, - None, - max_polls=0, - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def _start_and_test_with_notification_fallback_unlocked( - self, notification: bytes | None - ) -> bool: - """Keep tagged fallback start/observe/finalize under one slot lock.""" - - self._start_with_notification_unlocked(notification) - expected_generation = self._generation - state = self._refresh_unlocked() - if state is WorkState.COMPLETED: - self._finish_execute_success_unlocked(expected_generation) - return True - if state in _TERMINAL_STATES: - self._raise_for_state() - return False - - def execute( - self, - *, - notification: bytes | bytearray | memoryview | None, - timeout: float | None = None, - ) -> int: - """Execute and logically recycle one explicitly tagged generation.""" - - # Keep the common bytes/None cases inline: unlike the fixed slot this - # mode necessarily examines a request-local value, but it should not - # add a helper frame before the one fused pybind call. - if notification is not None: - if type(notification) is bytes: - if not notification: - raise ValueError( - "an attached notification payload must be non-empty" - ) - if len(notification) > _MAX_NOTIFICATION_BYTES: - raise ValueError( - "notification exceeds the " - f"{_MAX_NOTIFICATION_BYTES}-byte size limit" - ) - else: - notification = _validate_request_slot_notification(notification) - - fast_execute = self._notification_fast_execute - if timeout is None and fast_execute is not None: - next_generation = self._generation + 1 - try: - if self._closed: - raise ClosedError("transfer slot is closed") - if self._creation_failed: - raise BackendFailureError( - "transfer slot is retained for cleanup only" - ) - if ( - self._owner_thread is not None - and _GET_IDENT() != self._owner_thread - ): - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - if self._active: - raise BusyError( - "transfer slot already owns an unrecycled generation" - ) - - if not self._caller_ready: - ready_event_synchronize = self._ready_event_synchronize - if ready_event_synchronize is not None: - try: - ready_event_synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to wait for CUDA readiness" - ) from error - else: - bound = self._bound - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - try: - self._generation = next_generation - self._active = True - except BaseException: - self._active = False - self._generation = next_generation - 1 - self._state = None - self._error = None - self._start_error = None - self._backend_generation_active = False - raise - - state = fast_execute( - notification, - max_polls=_FUSED_REQUEST_SLOT_MAX_POLLS, - timeout_ns=None, - ) - except BaseException as error: - if self._generation != next_generation: - raise - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - if state is _WORK_COMPLETED: - self._active = False - return next_generation - if state is _WORK_RUNNING: - # Preserve the same ownership-before-state publication as - # the fixed slot. Disable the request-local dispatcher - # before the inherited continuation: an interruption after - # its idle commit must not leave a stale fast dispatcher. - self._backend_generation_active = True - self._state = state - self._notification_fast_execute = None - poll_bounded = self._backend_poll_bounded - if poll_bounded is not None: - return self._wait_execute_bounded_until( - None, - expected_generation=next_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - None, expected_generation=next_generation - ) - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot notification start returned " - "without acknowledging a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - elif not isinstance(state, WorkState): - self._start_error = BackendFailureError( - "backend returned invalid fused request-slot state " - f"{state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - state = refreshed - else: - state = self._observe_backend_state_unlocked(state) - - expected_generation = self._generation - if state in _TERMINAL_STATES: - self._raise_for_state() - # A bounded continuation finalizes through the inherited idle - # transaction. Disable this subclass's exact dispatcher first so - # an interruption after that transaction's public idle commit - # cannot pair it with a stale terminal state. - self._notification_fast_execute = None - poll_bounded = self._backend_poll_bounded - if poll_bounded is not None: - return self._wait_execute_bounded_until( - None, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - None, - expected_generation=expected_generation, - ) - - deadline_ns = None if timeout is None else _transfer_slot_deadline_ns(timeout) - start_and_poll = self._backend_start_and_poll_with_notification - lock = self._lock - poll_bounded = self._backend_poll_bounded - if start_and_poll is None: - if lock is None: - self._start_with_notification_unlocked(notification) - expected_generation = self._generation - else: - with lock: - self._start_with_notification_unlocked(notification) - expected_generation = self._generation - if poll_bounded is not None: - return self._wait_execute_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - deadline_ns, - expected_generation=expected_generation, - ) - - if lock is None: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, - notification, - deadline_ns, - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self._finish_execute_success_unlocked(expected_generation) - if state in _TERMINAL_STATES: - self._raise_for_state() - else: - with lock: - state = self._start_and_poll_with_notification_unlocked( - start_and_poll, - notification, - deadline_ns, - ) - expected_generation = self._generation - if state is WorkState.COMPLETED: - return self._finish_execute_success_unlocked(expected_generation) - if state in _TERMINAL_STATES: - self._raise_for_state() - - if deadline_ns is not None and time.monotonic_ns() >= deadline_ns: - self._validate_wait_generation(expected_generation) - raise TransferTimeoutError("transfer slot timed out") - if poll_bounded is not None: - return self._wait_execute_bounded_until( - deadline_ns, - expected_generation=expected_generation, - poll_bounded=poll_bounded, - ) - return self._wait_execute_until( - deadline_ns, - expected_generation=expected_generation, - ) - - def _start_with_notification_unlocked(self, notification: bytes | None) -> None: - previous_generation = self._generation - try: - self._begin_generation_unlocked() - self._notification_fast_execute = None - start = self._backend_start_with_notification - if start is None: - self._start_error = BackendFailureError( - "transfer slot has no notification-override operation" - ) - return - if self._lock is None: - result = start(notification) - else: - result = self._endpoint._backend_request_slot_start_with_notification( - start, notification - ) - self._backend_generation_active = True - if result is not None: - raise BackendFailureError( - "backend request-slot start_with_notification must return None" - ) - except BaseException as error: - if self._generation == previous_generation: - raise - self._start_error = error - - def _start_and_poll_with_notification_unlocked( - self, - start_and_poll: Callable[..., WorkState | None], - notification: bytes | None, - deadline_ns: int | None, - *, - max_polls: int = _FUSED_REQUEST_SLOT_MAX_POLLS, - ) -> WorkState: - previous_generation = self._generation - try: - self._begin_generation_unlocked() - self._notification_fast_execute = None - timeout_ns = ( - None - if deadline_ns is None - else min( - _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS, - max(0, deadline_ns - time.monotonic_ns()), - ) - ) - if self._lock is None: - state = start_and_poll( - notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - else: - state = self._endpoint._backend_request_slot_start_and_poll_with_notification( - start_and_poll, - notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - except BaseException as error: - if self._generation == previous_generation: - raise - self._start_error = error - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - if state is None: - self._start_error = BackendFailureError( - "backend fused request-slot notification start returned without " - "acknowledging a generation" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - if not isinstance(state, WorkState): - self._start_error = BackendFailureError( - f"backend returned invalid fused request-slot state {state!r}" - ) - refreshed = self._refresh_unlocked() - assert refreshed is not None - return refreshed - return self._observe_backend_state_unlocked(state) - - def _recycle_unlocked(self) -> None: - """Disable notification fast dispatch before an exceptional recycle.""" - - # Do not penalize or mutate an erroneous idle recycle. Active failure - # and cancellation recovery, however, may publish public idle before - # clearing retained diagnostics in the base transaction. Keeping the - # notification dispatcher disabled makes that dirty-idle boundary a - # safe generic fallback. - if self._active: - self._notification_fast_execute = None - super()._recycle_unlocked() - - def _restore_fast_execute_unlocked(self) -> None: - """Restore only the notification-aware exact-execute callable.""" - - self._fast_execute = None - start_and_poll = self._backend_start_and_poll_with_notification - self._notification_fast_execute = ( - start_and_poll - if ( - self._lock is None - and start_and_poll is not None - and self._backend_deferred_successful_recycle - ) - else None - ) - - def _adopt_backend_slot_unlocked( - self, - backend_slot: BackendRequestSlot, - *, - execution_dispatch: object | None, - ) -> None: - dispatch = backend_slot if execution_dispatch is None else execution_dispatch - start = getattr(dispatch, "start_with_notification", None) - start_and_poll = ( - getattr(dispatch, "start_and_poll_with_notification", None) - if self._endpoint.capabilities.fused_request_slot_start_poll - else None - ) - if not callable(start): - raise BackendFailureError( - "backend request-slot start_with_notification operation must be " - "callable" - ) - if ( - getattr(start, "__func__", None) - is BackendRequestSlot.start_with_notification - ): - raise BackendFailureError( - "backend advertised request-slot notification overrides without " - "implementing start_with_notification" - ) - if self._endpoint.capabilities.fused_request_slot_start_poll and not callable( - start_and_poll - ): - raise BackendFailureError( - "backend request-slot start_and_poll_with_notification operation " - "must be callable" - ) - if ( - start_and_poll is not None - and getattr(start_and_poll, "__func__", None) - is BackendRequestSlot.start_and_poll_with_notification - ): - raise BackendFailureError( - "backend advertised fused request-slot notification overrides " - "without implementing start_and_poll_with_notification" - ) - - # Publish the extra dispatch table before the base adoption commit. - # An interruption can therefore expose either no backend slot or one - # with every method required by its public dynamic type. - self._backend_start_with_notification = start - self._backend_start_and_poll_with_notification = start_and_poll - super()._adopt_backend_slot_unlocked( - backend_slot, - execution_dispatch=execution_dispatch, - ) - # Even an explicit unbound TransferSlot.start(dynamic_slot) must not - # enter the fixed provider family: NIXL's legacy nullptr options mean - # "retain the last tag", not "clear". Leave only observation, - # recycle, cancellation, and close shared between the two families. - self._backend_start = None - self._backend_start_and_poll = None - self._fast_execute = None - self._restore_fast_execute_unlocked() - - -class Work: - """Request-local asynchronous transfer state.""" - - __slots__ = ( - "_backend_work", - "_backend_released", - "_bound", - "_closed", - "_dependencies", - "_endpoint", - "_error", - "_owned_plan", - "_plan", - "_recovery_token", - "_lock", - "_owner_thread", - "_submission_error", - "_submitting", - "_state", - ) - - def __init__( - self, - endpoint: Endpoint, - backend_work: BackendWork | None, - *, - plan: TransferPlan | None = None, - bound: BoundTransfer | None = None, - dependencies: tuple[object, ...] = (), - owned_plan: bool = False, - recovery_token: object | None = None, - _token: object, - ) -> None: - if _token is not _CONSTRUCTION_TOKEN: - raise TypeError("Work objects are created by Endpoint operations") - self._endpoint = endpoint - self._backend_work = backend_work - self._backend_released = False - self._plan = plan - self._bound = bound - self._dependencies = dependencies - self._owned_plan = owned_plan - # Passive caller identity: never interpret or provider-forward it. - self._recovery_token = recovery_token - self._lock = ( - None - if endpoint._thread_mode - in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) - else threading.RLock() - ) - self._owner_thread = ( - endpoint._owner_thread - if endpoint._thread_mode is ThreadMode.SINGLE - else None - ) - self._submission_error: BaseException | None = None - # A Work is published before provider entry so an interrupted caller can - # recover its identity. Observation is unsafe until submission return or - # exception reconciliation has completed. - self._submitting = True - self._state = WorkState.PENDING - self._error: BaseException | None = None - self._closed = False - - @property - def state(self) -> WorkState: - return self._refresh() - - @property - def recovery_token(self) -> object | None: - """Return the opaque identity retained for lost-return recovery. - - Match this object with ``is`` in :attr:`Endpoint.open_works`. Core - never hashes, compares, invokes, serializes, or passes it to a provider. - """ - - return self._recovery_token - - @property - def error(self) -> BaseException | None: - self._refresh() - return self._error - - def exception(self) -> BaseException | None: - return self.error - - @property - def completion_event(self) -> object | None: - lock = self._lock - if lock is None: - return self._completion_event_unlocked() - with lock: - return self._completion_event_unlocked() - - def _completion_event_unlocked(self) -> object | None: - self._check_not_closed() - self._check_submission_complete_unlocked() - if not self._endpoint.capabilities.cuda_completion_event: - return None - if self._backend_work is None: - raise BusyError("transfer work has not been adopted by its backend") - return self._endpoint._backend_call( - "query work completion event", - lambda: self._backend_work.completion_event, - ) - - def _adopt_backend_work(self, backend_work: BackendWork) -> None: - """Adopt provider state without allocating a submission closure.""" - - if not isinstance(backend_work, BackendWork): - raise BackendFailureError("backend must adopt a BackendWork instance") - lock = self._lock - if lock is None: - self._adopt_backend_work_unlocked(backend_work) - return - with lock: - self._adopt_backend_work_unlocked(backend_work) - - def _adopt_backend_work_unlocked(self, backend_work: BackendWork) -> None: - if self._closed: - raise BackendFailureError( - "backend attempted to adopt an expired transfer work" - ) - if self._backend_work is not None: - raise BackendFailureError( - "backend attempted to adopt transfer work more than once" - ) - self._backend_work = backend_work - - @property - def closed(self) -> bool: - return self._closed - - def test(self) -> bool: - return self._refresh() in _TERMINAL_STATES - - def poll_bounded( - self, - *, - max_polls: int = _FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - timeout_ns: int | None = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, - timeout_check_interval: int = (_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL), - ) -> WorkState: - """Perform a bounded poll and return the exact observed state. - - A provider advertising ``fused_work_poll`` performs the entire loop in - one provider call. Other providers retain identical semantics through - repeated ordinary state observations. The returned state may be - nonterminal; this Work remains the authoritative owner. - """ - - _validate_bounded_poll_arguments( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - allow_zero=False, - ) - if self._endpoint.capabilities.fused_work_poll: - lock = self._lock - if lock is None: - return self._poll_bounded_unlocked( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - with lock: - return self._poll_bounded_unlocked( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - - deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns - state = self._state - for poll_index in range(max_polls): - state = self._refresh() - if state in _TERMINAL_STATES: - return state - if ( - deadline_ns is not None - and (poll_index == 0 or (poll_index + 1) % timeout_check_interval == 0) - and time.monotonic_ns() >= deadline_ns - ): - break - return state - - def _poll_bounded_unlocked( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState: - self._check_not_closed() - self._check_thread_unlocked() - self._check_submission_complete_unlocked() - if self._state in _TERMINAL_STATES: - return self._state - backend_work = self._backend_work - if backend_work is None: - return self._state - state = self._endpoint._backend_work_poll_bounded( - backend_work, - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - self._ingest_backend_state_unlocked(state) - return self._state - - def test_bounded( - self, - *, - max_polls: int = _FUSED_REQUEST_SLOT_POLL_MAX_POLLS, - timeout_ns: int | None = _FUSED_REQUEST_SLOT_POLL_TIMEOUT_NS, - timeout_check_interval: int = (_FUSED_REQUEST_SLOT_POLL_TIMEOUT_CHECK_INTERVAL), - ) -> bool: - """Compatibility wrapper returning whether a bounded poll is terminal.""" - - return ( - self.poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - in _TERMINAL_STATES - ) - - def wait(self, timeout: float | None = None) -> Work: - if timeout is not None: - if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): - raise TypeError("timeout must be a non-negative number or None") - if not math.isfinite(timeout) or timeout < 0: - raise ValueError("timeout must be finite and non-negative") - deadline = None if timeout is None else time.monotonic() + timeout - polls = 0 - while True: - state = self._refresh() - if state in _TERMINAL_STATES: - return self._raise_for_state() - if deadline is not None and time.monotonic() >= deadline: - raise TransferTimeoutError("transfer work timed out") - if self._endpoint.progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - polls += 1 - if polls >= 16: - # Yield without imposing a fixed millisecond latency floor. - # Long waits back off to cap host polling overhead. - time.sleep(0 if polls < 128 else 0.0001) - - def result(self, timeout: float | None = None) -> None: - self.wait(timeout) - return None - - def cancel(self) -> bool: - lock = self._lock - if lock is None: - return self._cancel_unlocked() - with lock: - return self._cancel_unlocked() - - def _cancel_unlocked(self) -> bool: - self._check_not_closed() - self._check_submission_complete_unlocked() - if self._refresh_unlocked() in _TERMINAL_STATES: - return self._state is WorkState.CANCELLED - if self._backend_work is None: - return False - cancelled = self._endpoint._backend_call( - "cancel work", self._backend_work.cancel - ) - if not isinstance(cancelled, bool): - raise BackendFailureError("backend work cancellation must return a bool") - self._refresh_unlocked() - return cancelled and self._state is WorkState.CANCELLED - - def close(self) -> None: - state_lock = self._endpoint._state_lock - if state_lock is _NOOP_LOCK: - work_lock = self._lock - if work_lock is None: - self._close_unlocked() - return - with work_lock: - self._close_unlocked() - return - with state_lock: - self._close_with_work_lock() - - def _close_with_work_lock(self) -> None: - lock = self._lock - if lock is None: - self._close_unlocked() - return - with lock: - self._close_unlocked() - - def _close_unlocked(self) -> None: - if self._closed: - # CLOSED is the provider-release commit, not proof that every - # Core-only graph edge and retained reference was cleared. A - # retry must finish an epilogue interrupted after that commit. - self._finish_close_unlocked() - return - self._check_submission_complete_unlocked() - self._check_thread_unlocked() - if ( - self._state not in _TERMINAL_STATES - and self._refresh_unlocked() not in _TERMINAL_STATES - ): - raise BusyError("cannot close non-terminal transfer work") - if not self._backend_released: - if self._backend_work is None: - raise BackendFailureError( - "terminal transfer work has no backend request" - ) - self._endpoint._backend_work_close(self._backend_work) - self._backend_released = True - if self._plan is not None: - self._plan._works.discard(self) - if self._bound is not None: - self._bound._works.discard(self) - try: - if self._owned_plan and self._plan is not None: - self._plan.close() - except BaseException: - if self._plan is not None: - self._plan._works.add(self) - raise - self._closed = True - self._finish_close_unlocked() - - def _finish_close_unlocked(self) -> None: - """Idempotently unlink a closed Work from the Core lifetime graph.""" - - plan = self._plan - bound = self._bound - # Retain both parent references until every blocking edge is gone. An - # interrupt during any discard therefore leaves enough state to retry. - if plan is not None: - plan._works.discard(self) - if bound is not None: - bound._works.discard(self) - self._endpoint._works.discard(self) - self._dependencies = () - self._plan = None - self._bound = None - self._backend_work = None - self._error = None - self._submission_error = None - # Commit last: absence means every provider/Core lifetime edge is gone. - self._recovery_token = None - - async def _wait_async(self) -> Work: - while True: - state = self._refresh() - if state in _TERMINAL_STATES: - return self._raise_for_state() - if self._endpoint.progress_mode is ProgressMode.MANUAL: - self._endpoint.progress() - await asyncio.sleep(0) - - def __await__(self): # type: ignore[no-untyped-def] - return self._wait_async().__await__() - - def _refresh(self) -> WorkState: - lock = self._lock - if lock is None: - return self._refresh_unlocked() - with lock: - return self._refresh_unlocked() - - def _refresh_unlocked(self) -> WorkState: - self._check_not_closed() - self._check_thread_unlocked() - self._check_submission_complete_unlocked() - if self._state in _TERMINAL_STATES: - return self._state - if self._backend_work is None: - return self._state - state = self._endpoint._backend_work_state(self._backend_work) - return self._ingest_backend_state_unlocked(state) - - def _ingest_backend_state_unlocked(self, state: WorkState) -> WorkState: - """Validate and cache a provider state already observed this call.""" - - if not isinstance(state, WorkState): - raise BackendFailureError(f"backend returned invalid work state {state!r}") - if self._state is WorkState.RUNNING and state is WorkState.PENDING: - raise BackendFailureError( - "backend work state regressed from running to pending" - ) - observed_state = state - observed_error = self._error - if state in _TERMINAL_STATES and self._submission_error is not None: - observed_state = WorkState.FAILED - observed_error = self._submission_error - elif state is WorkState.FAILED: - backend_work = self._backend_work - if backend_work is None: - observed_error = BackendFailureError( - "failed transfer work has no backend request" - ) - else: - try: - error = self._endpoint._backend_call( - "query work error", lambda: backend_work.error - ) - except Exception as observation_error: - observed_error = observation_error - else: - if error is None: - error = BackendFailureError( - "transfer backend reported failure without an error" - ) - elif not isinstance(error, BaseException): - error = BackendFailureError( - "transfer backend returned a non-exception error" - ) - observed_error = error - elif state is WorkState.CANCELLED: - observed_error = TransferCancelledError("transfer work was cancelled") - - # Terminal state is the public commit. Derive its diagnostic first so - # an asynchronous BaseException during provider error observation leaves - # the prior nonterminal state retryable. - if observed_state in _TERMINAL_STATES: - self._error = observed_error - self._state = observed_state - return self._state - - def _raise_for_state(self) -> Work: - if self._state is WorkState.COMPLETED: - return self - if self._error is None: - self._error = BackendFailureError( - f"transfer work reached {self._state.value} without an error" - ) - if isinstance(self._error, TransferError): - raise self._error - raise BackendFailureError("transfer work failed") from self._error - - def _check_not_closed(self) -> None: - if self._closed: - raise ClosedError("transfer work is closed") - - def _check_submission_complete_unlocked(self) -> None: - if self._submitting: - raise BusyError("transfer work submission is in progress") - - def _check_thread_unlocked(self) -> None: - owner_thread = self._owner_thread - if owner_thread is not None and _GET_IDENT() != owner_thread: - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - - -class Endpoint: - """Backend-neutral endpoint for asymmetric one-sided data transfer.""" - - def __init__( - self, - name: str, - *, - backend: str = "reference", - endpoint_id: str | None = None, - progress_mode: ProgressMode | str = ProgressMode.MANUAL, - thread_mode: ThreadMode | str = ThreadMode.SERIALIZED, - options: Mapping[str, Any] | None = None, - ) -> None: - _validate_nonempty_str(name, "name") - _validate_nonempty_str(backend, "backend") - if endpoint_id is not None: - _validate_nonempty_str(endpoint_id, "endpoint_id") - progress_mode = _coerce_enum(progress_mode, ProgressMode, "progress_mode") - thread_mode = _coerce_enum(thread_mode, ThreadMode, "thread_mode") - if options is None: - options = {} - elif not isinstance(options, Mapping): - raise TypeError("options must be a mapping or None") - - self._name = name - self._backend_name = backend - self._endpoint_id = endpoint_id or str(uuid.uuid4()) - self._incarnation = str(uuid.uuid4()) - self._progress_mode = progress_mode - self._thread_mode = thread_mode - self._owner_thread = threading.get_ident() - self._backend_lock = ( - threading.RLock() if thread_mode is ThreadMode.SERIALIZED else _NOOP_LOCK - ) - self._state_lock = ( - _NOOP_LOCK - if thread_mode in (ThreadMode.SINGLE, ThreadMode.CALLER_SERIALIZED) - else threading.RLock() - ) - self._notification_consumer_lock = ( - threading.RLock() if thread_mode is ThreadMode.MULTIPLE else _NOOP_LOCK - ) - self._active_multiple_backend_calls: ( - set[weakref.ReferenceType[_BackendCallReservation]] | None - ) = (set() if thread_mode is ThreadMode.MULTIPLE else None) - self._multiple_backend_call_reservation_pool: ( - list[_BackendCallReservation] | None - ) = ([] if thread_mode is ThreadMode.MULTIPLE else None) - self._registrations: dict[str, Registration] = {} - self._registration_names: dict[str, Registration] = {} - self._peers: dict[str, Peer] = {} - self._plans: set[TransferPlan] = set() - self._orphaned_plans: set[TransferPlan] = set() - self._bound_transfers: set[BoundTransfer] = set() - self._transfer_slots: set[TransferSlot] = set() - self._works: set[Work] = set() - self._closing = False - self._creation_failed = False - self._closed = False - - _retry_orphaned_backend_receipts(backend, self._endpoint_id) - registration = _get_backend_registration(backend) - backend_options = MappingProxyType(dict(options)) - receipt = _BackendConstructionReceipt(backend, self._endpoint_id) - self._backend_receipt = receipt - self._backend_finalizer = weakref.finalize( - self, _finalize_backend_receipt, receipt - ) - self._backend_finalizer.atexit = False - endpoint_ref = weakref.ref(self) - - def adopt_backend(implementation: Backend) -> None: - # The receipt owns the backend before an Endpoint attribute becomes - # observable. The outer transaction repairs either STORE_ATTR cut. - receipt.adopt(implementation) - endpoint = endpoint_ref() - if endpoint is None: - raise BackendFailureError( - "endpoint disappeared while adopting its backend" - ) - endpoint._backend = implementation - endpoint._creation_failed = True - - factory = registration.factory - unowned_v1_result = False - try: - if registration.factory_api_version == 1: - candidate = cast(Callable[..., object], factory)( - self._endpoint_id, - name, - self._incarnation, - progress_mode, - thread_mode, - backend_options, - ) - if not isinstance(candidate, Backend): - # A wrong-type legacy return was never adopted by Core. It - # may be an arbitrary foreign object, so do not call close - # or inspect any provider-shaped attribute on it. - unowned_v1_result = True - raise TypeError("backend factory must return a Backend instance") - # This compatibility path necessarily retains the legacy - # CALL-to-adoption gap. Version-2 providers close it. - adopt_backend(candidate) - implementation = receipt.seal_successful_call(None) - else: - result = cast(Callable[..., object], factory)( - self._endpoint_id, - name, - self._incarnation, - progress_mode, - thread_mode, - backend_options, - adopt_backend=adopt_backend, - ) - implementation = receipt.seal_successful_call(result) - except BaseException as error: - receipt.seal_failed_call() - try: - self._cleanup_adopted_backend_after_creation_error(error) - except BaseException: - raise - if unowned_v1_result: - raise - if isinstance(error, TransferError) or not isinstance(error, Exception): - raise - raise BackendFailureError( - f"failed to create transfer backend {backend!r}" - ) from error - - try: - # From this assignment onward Core owns the returned backend. Keep - # capability discovery, dispatch caching, and EndpointInfo - # publication in one cleanup transaction. - self._backend = implementation - self._creation_failed = True - capabilities = implementation.capabilities - if not isinstance(capabilities, Capabilities): - raise TypeError("backend capabilities must be Capabilities") - if ( - progress_mode is ProgressMode.MANUAL - and not capabilities.manual_progress - ): - raise UnsupportedError( - f"backend {backend!r} does not support manual progress" - ) - if ( - progress_mode is ProgressMode.BACKGROUND - and not capabilities.background_progress - ): - raise UnsupportedError( - f"backend {backend!r} does not support background progress" - ) - supported_thread_modes = implementation.supported_thread_modes - if not isinstance(supported_thread_modes, frozenset) or not all( - isinstance(mode, ThreadMode) for mode in supported_thread_modes - ): - raise TypeError( - "backend supported_thread_modes must be a frozenset of " - "ThreadMode values" - ) - if thread_mode not in supported_thread_modes: - raise UnsupportedError( - f"backend {backend!r} does not support thread mode " - f"{thread_mode.value!r}" - ) - send_notification_sync_dispatch = ( - cast( - Callable[[object, bytes], object | None], - implementation.send_notification_sync, - ) - if capabilities.synchronous_notification_send - else None - ) - info = EndpointInfo( - backend=backend, - endpoint_id=self._endpoint_id, - name=name, - incarnation=self._incarnation, - capabilities=capabilities, - ) - self._capabilities = capabilities - self._send_notification_sync_dispatch = send_notification_sync_dispatch - self._info = info - self._creation_failed = False - receipt.mark_published() - except BaseException as error: - # Re-establish both fields if interruption followed either - # STORE_ATTR. Failed cleanup is exposed through a recovery endpoint. - self._backend = implementation - self._creation_failed = True - try: - self._close_failed_endpoint_backend() - except BaseException as cleanup_error: - raise _endpoint_creation_cleanup_failure( - self, error, cleanup_error - ) from cleanup_error - raise - - @property - def name(self) -> str: - return self._name - - @property - def backend(self) -> str: - return self._backend_name - - @property - def endpoint_id(self) -> str: - return self._endpoint_id - - @property - def incarnation(self) -> str: - return self._incarnation - - @property - def identity(self) -> EndpointIdentity: - return EndpointIdentity(self.endpoint_id, self.incarnation) - - @property - def progress_mode(self) -> ProgressMode: - return self._progress_mode - - @property - def thread_mode(self) -> ThreadMode: - return self._thread_mode - - @property - def capabilities(self) -> Capabilities: - return self._capabilities - - @property - def info(self) -> EndpointInfo: - return self._info - - @property - def closed(self) -> bool: - return self._closed - - @property - def open_registrations(self) -> tuple[Registration, ...]: - """Return registrations owned by this endpoint for recovery/inspection.""" - - with self._state_lock: - self._check_open() - return tuple(self._registrations.values()) - - @property - def open_peers(self) -> tuple[Peer, ...]: - """Return imported peers owned by this endpoint for recovery/inspection.""" - - with self._state_lock: - self._check_open() - return tuple(self._peers.values()) - - @property - def open_plans(self) -> tuple[TransferPlan, ...]: - """Return plans owned by this endpoint for recovery/inspection.""" - - with self._state_lock: - self._check_open() - return tuple(self._plans) - - @property - def open_works(self) -> tuple[Work, ...]: - """Return open Work objects, including interrupted submissions. - - The snapshot lets an application recover and drain a Work if an - asynchronous exception lands after submission but before the caller - stores the returned object. - """ - - with self._state_lock: - self._check_open() - return tuple(self._works) - - @property - def open_bound_transfers(self) -> tuple[BoundTransfer, ...]: - """Return reusable bindings owned by this endpoint for recovery.""" - - with self._state_lock: - self._check_open() - return tuple(self._bound_transfers) - - @property - def open_transfer_slots(self) -> tuple[TransferSlot, ...]: - """Return persistent request slots owned by this endpoint for recovery.""" - - with self._state_lock: - self._check_open() - return tuple(self._transfer_slots) - - @_state_locked - def register( - self, - tensor_or_storage: object, - *, - name: str | None = None, - ) -> Registration: - self._check_open() - if name is None: - name = f"registration-{uuid.uuid4()}" - _validate_nonempty_str(name, "name") - if name in self._registration_names: - raise InvalidRegionError(f"registration name {name!r} is not unique") - registration_id = str(uuid.uuid4()) - memory = _registered_memory(tensor_or_storage, registration_id, name) - for existing in self._registrations.values(): - _validate_memory_stable(existing._memory) - if _memory_ranges_overlap(memory, existing._memory): - raise InvalidRegionError( - f"registration {name!r} overlaps existing registration " - f"{existing.name!r}" - ) - if isinstance(tensor_or_storage, RawSpan) and not self.capabilities.raw_spans: - raise UnsupportedError( - f"backend {self.backend!r} does not support raw spans" - ) - registration = Registration( - self, - memory, - _UNADOPTED_HANDLE, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._registrations[registration_id] = registration - self._registration_names[name] = registration - result = self._backend_call( - "register memory", - lambda: self._backend.register( - memory, - adopt_handle=lambda handle: _adopt_resource_handle( - registration, handle, "registration" - ), - ), - ) - if result is not None: - raise BackendFailureError( - "backend register must return None after adopting its handle" - ) - if registration._handle is _UNADOPTED_HANDLE: - raise BackendFailureError( - "backend register returned without adopting a handle" - ) - return registration - except BaseException as error: - if registration._handle is not _UNADOPTED_HANDLE: - try: - registration.close() - except BaseException as cleanup_error: - registration._creation_failed = True - raise _creation_cleanup_failure( - "registration", error, cleanup_error - ) from cleanup_error - else: - self._registrations.pop(registration_id, None) - self._registration_names.pop(name, None) - registration._closed = True - registration._handle = None - registration._memory = _released_memory(memory) - raise - - @_state_locked - def export_metadata( - self, - registrations: Sequence[Registration] | None = None, - ) -> EndpointMetadata: - self._check_open() - selected = ( - tuple(self._registrations.values()) - if registrations is None - else tuple(registrations) - ) - if len(selected) != len(set(selected)): - raise InvalidMetadataError("registrations must not contain duplicates") - values: list[RegistrationMetadata] = [] - for registration in selected: - self._validate_registration(registration) - descriptor = self._backend_call( - "describe registration", - lambda registration=registration: ( - self._backend.describe_registration(registration._handle) - ), - ) - if not isinstance(descriptor, (bytes, bytearray, memoryview)): - raise BackendFailureError( - "backend registration descriptor must be bytes-like" - ) - values.append( - RegistrationMetadata( - registration_id=registration.registration_id, - name=registration.name, - nbytes=registration.nbytes, - device=registration.device, - memory_type=registration.memory_type, - backend_descriptor=bytes(descriptor), - ) - ) - payload = self._backend_call( - "export endpoint metadata", - lambda: self._backend.endpoint_payload( - tuple(registration._handle for registration in selected) - ), - ) - if not isinstance(payload, (bytes, bytearray, memoryview)): - raise BackendFailureError("backend endpoint payload must be bytes-like") - return EndpointMetadata( - wire_major=EndpointMetadata.WIRE_MAJOR, - wire_minor=EndpointMetadata.WIRE_MINOR, - required_features=(), - backend=self.backend, - endpoint_id=self.endpoint_id, - name=self.name, - incarnation=self.incarnation, - backend_payload=bytes(payload), - registrations=tuple(values), - ) - - @_state_locked - def import_peer( - self, - metadata: EndpointMetadata | bytes | bytearray | memoryview, - *, - expected_identity: EndpointIdentity | None = None, - expected_endpoint_id: str | None = None, - expected_name: str | None = None, - expected_incarnation: str | None = None, - ) -> Peer: - self._check_open() - if not isinstance(metadata, EndpointMetadata): - metadata = EndpointMetadata.from_bytes(metadata) - if expected_identity is not None: - if not isinstance(expected_identity, EndpointIdentity): - raise TypeError("expected_identity must be an EndpointIdentity") - if expected_endpoint_id is not None or expected_incarnation is not None: - raise ValueError( - "expected_identity is mutually exclusive with separate expected " - "endpoint_id/incarnation values" - ) - expected_endpoint_id = expected_identity.endpoint_id - expected_incarnation = expected_identity.incarnation - expected_fields = { - "endpoint_id": expected_endpoint_id, - "name": expected_name, - "incarnation": expected_incarnation, - } - for field_name, expected in expected_fields.items(): - if expected is None: - continue - _validate_nonempty_str(expected, f"expected_{field_name}") - actual = getattr(metadata, field_name) - if actual != expected: - raise InvalidMetadataError( - f"peer {field_name} {actual!r} does not match expected " - f"value {expected!r}" - ) - if metadata.backend != self.backend: - raise BackendMismatchError( - f"peer backend {metadata.backend!r} does not match " - f"local backend {self.backend!r}" - ) - existing = self._peers.get(metadata.endpoint_id) - if existing is not None: - if existing.incarnation != metadata.incarnation: - raise BusyError( - f"peer {metadata.endpoint_id!r} changed incarnation; " - "close the old peer before importing the replacement" - ) - raise BusyError(f"peer {metadata.endpoint_id!r} is already imported") - peer = Peer( - self, - metadata, - _UNADOPTED_HANDLE, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._peers[peer.endpoint_id] = peer - result = self._backend_call( - "import peer", - lambda: self._backend.import_peer( - metadata, - adopt_handle=lambda handle: _adopt_resource_handle( - peer, handle, "peer" - ), - ), - ) - if result is not None: - raise BackendFailureError( - "backend peer import must return None after adopting its handle" - ) - if peer._handle is _UNADOPTED_HANDLE: - raise BackendFailureError( - "backend peer import returned without adopting a handle" - ) - return peer - except BaseException as error: - if peer._handle is not _UNADOPTED_HANDLE: - try: - peer.close() - except BaseException as cleanup_error: - peer._creation_failed = True - raise _creation_cleanup_failure( - "peer", error, cleanup_error - ) from cleanup_error - else: - self._peers.pop(peer.endpoint_id, None) - peer._closed = True - peer._handle = None - peer._registrations = {} - raise - - @_state_locked - def prepare( - self, - *, - local: Sequence[Region], - remote: Sequence[Region], - indexed: bool = False, - ) -> TransferPlan: - self._check_open() - if not isinstance(indexed, bool): - raise TypeError("indexed must be a bool") - local = tuple(local) - remote = tuple(remote) - self._validate_regions(local, remote, indexed=indexed) - backend_local = tuple(region._to_backend() for region in local) - backend_remote = tuple(region._to_backend() for region in remote) - # Finish every Python allocation and publish the local lifetime guard - # before asking a provider to create native state. - plan = TransferPlan( - self, - _UNADOPTED_HANDLE, - local, - remote, - indexed, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._plans.add(plan) - result = self._backend_call( - "prepare transfer plan", - lambda: self._backend.prepare( - backend_local, - backend_remote, - indexed=indexed, - adopt_handle=lambda handle: _adopt_resource_handle( - plan, handle, "transfer plan" - ), - ), - ) - if result is not None: - raise BackendFailureError( - "backend prepare must return None after adopting its handle" - ) - if plan._handle is _UNADOPTED_HANDLE: - raise BackendFailureError( - "backend prepare returned without adopting a handle" - ) - return plan - except BaseException as error: - if plan._handle is not _UNADOPTED_HANDLE: - try: - plan.close() - except BaseException as cleanup_error: - plan._creation_failed = True - self._orphaned_plans.add(plan) - raise _creation_cleanup_failure( - "transfer plan", error, cleanup_error - ) from cleanup_error - else: - self._plans.discard(plan) - plan._closed = True - plan._handle = None - plan._local = () - plan._remote = () - plan._origins = frozenset() - plan._reuse_tokens = {} - raise - - def read( - self, - plan: TransferPlan | None = None, - *, - local: Sequence[Region] | None = None, - remote: Sequence[Region] | None = None, - local_indices: Sequence[int] | None = None, - remote_indices: Sequence[int] | None = None, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - ) -> Work: - return self._submit( - TransferOp.READ, - plan, - local=local, - remote=remote, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - - def write( - self, - plan: TransferPlan | None = None, - *, - local: Sequence[Region] | None = None, - remote: Sequence[Region] | None = None, - local_indices: Sequence[int] | None = None, - remote_indices: Sequence[int] | None = None, - notification: bytes | bytearray | memoryview | None = None, - cuda_ordering: CudaOrderingMode | str = CudaOrderingMode.CALLER_READY, - stream: object | None = None, - ready_event: object | None = None, - recovery_token: object | None = None, - ) -> Work: - return self._submit( - TransferOp.WRITE, - plan, - local=local, - remote=remote, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - - @_state_locked - def _select_plan_indices( - self, - plan: TransferPlan, - op: TransferOp | str, - *, - local_indices: Sequence[int], - remote_indices: Sequence[int], - ) -> IndexSelection: - self._check_open() - self._validate_plan_locked(plan) - if not plan.indexed: - raise InvalidRegionError( - "indices require a plan prepared with indexed=True" - ) - op = _coerce_enum(op, TransferOp, "op") - supported = ( - self.capabilities.read if op is TransferOp.READ else self.capabilities.write - ) - if not supported: - raise UnsupportedError( - f"backend {self.backend!r} does not support {op.value}" - ) - backend_selection, local_devices = _snapshot_backend_index_selection( - plan, op, local_indices, remote_indices - ) - return IndexSelection( - self, - plan, - op, - backend_selection, - local_devices, - _token=_CONSTRUCTION_TOKEN, - ) - - @_state_locked - def _bind_transfer( - self, - plan: TransferPlan, - op: TransferOp | str, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | bytearray | memoryview | None, - cuda_ordering: CudaOrderingMode | str, - stream: object | None, - ready_event: object | None, - ) -> BoundTransfer: - self._check_open() - self._validate_plan_locked(plan) - op = _coerce_enum(op, TransferOp, "op") - supported = ( - self.capabilities.read if op is TransferOp.READ else self.capabilities.write - ) - if not supported: - raise UnsupportedError( - f"backend {self.backend!r} does not support {op.value}" - ) - cuda_ordering = _coerce_enum(cuda_ordering, CudaOrderingMode, "cuda_ordering") - _validate_cuda_ordering_arguments( - self.capabilities, - self.backend, - cuda_ordering, - stream, - ready_event, - ) - local_indices, remote_indices = _validate_plan_indices( - plan, local_indices, remote_indices, op - ) - notification = _validate_attached_notification( - self.capabilities, self.backend, notification - ) - cuda_device = _validate_cuda_ordering_selection( - plan, - local_indices, - cuda_ordering, - stream, - ready_event, - ) - binding = BoundTransfer( - self, - plan, - op, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - cuda_device=cuda_device, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - selection=None, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._bound_transfers.add(binding) - plan._bindings.add(binding) - except BaseException: - plan._bindings.discard(binding) - self._bound_transfers.discard(binding) - binding._closed = True - binding._plan = None - raise - return binding - - @_state_locked - def _bind_index_selection( - self, - selection: IndexSelection, - *, - cuda_ordering: CudaOrderingMode | str, - stream: object | None, - ready_event: object | None, - ) -> BoundTransfer: - self._check_open() - if type(selection) is not IndexSelection: - raise TypeError("expected an IndexSelection") - if selection._endpoint is not self: - raise InvalidRegionError("index selection belongs to another endpoint") - plan = selection._check_open() - cuda_ordering = _coerce_enum(cuda_ordering, CudaOrderingMode, "cuda_ordering") - _validate_cuda_ordering_arguments( - self.capabilities, - self.backend, - cuda_ordering, - stream, - ready_event, - ) - cuda_device = _validate_cuda_ordering_devices( - selection._local_devices, - cuda_ordering, - stream, - ready_event, - ) - binding = BoundTransfer( - self, - plan, - selection._op, - local_indices=selection._backend_selection.local_indices, - remote_indices=selection._backend_selection.remote_indices, - notification=None, - cuda_device=cuda_device, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - selection=selection, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._bound_transfers.add(binding) - plan._bindings.add(binding) - except BaseException: - plan._bindings.discard(binding) - self._bound_transfers.discard(binding) - binding._closed = True - binding._plan = None - binding._selection = None - raise - return binding - - @_state_locked - def _create_transfer_slot( - self, - bound: BoundTransfer, - *, - request_local_notifications: bool, - ) -> TransferSlot | RequestLocalNotificationSlot: - """Publish Core ownership before asking a provider for native state.""" - - self._check_open() - if not self._capabilities.persistent_request_slots: - raise UnsupportedError( - f"backend {self.backend!r} does not support persistent request slots" - ) - if not isinstance(bound, BoundTransfer): - raise TypeError("expected a BoundTransfer") - if bound._endpoint is not self: - raise InvalidRegionError("bound transfer belongs to another endpoint") - if request_local_notifications: - if not self._capabilities.request_slot_notification_overrides: - raise UnsupportedError( - f"backend {self.backend!r} does not support request-local " - "persistent-slot notifications" - ) - if bound._notification is not None: - raise ValueError( - "request-local notification slots require a binding with " - "notification=None" - ) - plan = bound._check_open() - slot: TransferSlot - if request_local_notifications: - slot = RequestLocalNotificationSlot(self, bound, _token=_CONSTRUCTION_TOKEN) - else: - slot = TransferSlot(self, bound, _token=_CONSTRUCTION_TOKEN) - slot._creation_failed = True - try: - # Publish the recovery registry first. At every subsequent - # interruption point either normal rollback succeeds or the slot - # remains discoverable through ``open_transfer_slots``. - self._transfer_slots.add(slot) - bound._slots.add(slot) - result = self._backend_create_request_slot( - bound._op, - plan._handle, - local_indices=bound._local_indices, - remote_indices=bound._remote_indices, - notification=bound._notification, - adopt_slot=slot._adopt_backend_slot, - ) - if result is not None: - raise BackendFailureError( - "backend request-slot creation must return None after adoption" - ) - if slot._backend_slot is None: - raise BackendFailureError( - "backend request-slot creation returned without adoption" - ) - slot._creation_failed = False - return slot - except BaseException as error: - if slot._backend_slot is not None: - try: - self._close_transfer_slot(slot) - except BaseException as cleanup_error: - slot._creation_failed = True - raise _creation_cleanup_failure( - "transfer slot", error, cleanup_error - ) from cleanup_error - else: - # Keep the endpoint recovery edge until the bound blocker has - # been removed. If cleanup itself is interrupted, the slot is - # still public through the endpoint snapshot and closeable. - bound._slots.discard(slot) - slot._closed = True - self._finish_transfer_slot_close(slot) - raise - - def _reserve_indexed_selection_submission( - self, - *, - plan: TransferPlan | None, - op: TransferOp | str | None, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - selection: IndexSelection | None, - notification: bytes | bytearray | memoryview | None, - cuda_ordering: CudaOrderingMode | str, - stream: object | None, - ready_event: object | None, - recovery_token: object | None, - ) -> tuple[ - TransferPlan, - TransferOp, - BackendIndexSelection, - bytes | None, - CudaOrderingMode, - object | None, - Work, - ]: - """Validate and construct an unpublished indexed Work under state lock. - - The caller holds the lock through publication. Returning a published - Work through a decorated helper would expose a lock-exit ownership gap. - """ - - active_plan: TransferPlan | None = None - work: Work | None = None - try: - self._check_open() - if selection is None: - if plan is None: - raise RuntimeError( - "one-shot indexed submission is missing its internal plan" - ) - self._validate_plan_locked(plan) - if not plan.indexed: - raise InvalidRegionError( - "indices require a plan prepared with indexed=True" - ) - active_plan = plan - active_op = _coerce_enum(op, TransferOp, "op") - supported = ( - self.capabilities.read - if active_op is TransferOp.READ - else self.capabilities.write - ) - if not supported: - raise UnsupportedError( - f"backend {self.backend!r} does not support {active_op.value}" - ) - caller_ready_hint = bool( - cuda_ordering is CudaOrderingMode.CALLER_READY - or ( - type(cuda_ordering) is str - and cuda_ordering == CudaOrderingMode.CALLER_READY.value - ) - ) - backend_selection, local_devices = _snapshot_backend_index_selection( - active_plan, - active_op, - local_indices, - remote_indices, - collect_local_devices=not caller_ready_hint, - ) - selection_dependency: object = backend_selection - else: - if type(selection) is not IndexSelection: - raise TypeError("expected an IndexSelection") - if selection._endpoint is not self: - raise InvalidRegionError( - "index selection belongs to another endpoint" - ) - active_plan = selection._plan - self._validate_plan_locked(active_plan) - active_op = selection._op - backend_selection = selection._backend_selection - local_devices = selection._local_devices - selection_dependency = selection - cuda_ordering = _coerce_enum( - cuda_ordering, CudaOrderingMode, "cuda_ordering" - ) - _validate_cuda_ordering_arguments( - self.capabilities, - self.backend, - cuda_ordering, - stream, - ready_event, - ) - notification = _validate_attached_notification( - self.capabilities, self.backend, notification - ) - cuda_device = ( - None - if cuda_ordering is CudaOrderingMode.CALLER_READY - else _validate_cuda_ordering_devices( - local_devices, - cuda_ordering, - stream, - ready_event, - ) - ) - if stream is None and ready_event is None: - dependencies = ( - (selection_dependency,) - if notification is None - else (selection_dependency, notification) - ) - else: - dependencies = tuple( - value - for value in ( - selection_dependency, - stream, - ready_event, - notification, - ) - if value is not None - ) - work = Work( - self, - None, - plan=active_plan, - dependencies=dependencies, - recovery_token=recovery_token, - _token=_CONSTRUCTION_TOKEN, - ) - return ( - active_plan, - active_op, - backend_selection, - notification, - cuda_ordering, - cuda_device, - work, - ) - except BaseException: - if work is not None: - work_lock = work._lock - if work_lock is None: - work._closed = True - work._submitting = False - work._dependencies = () - work._plan = None - work._recovery_token = None - else: - with work_lock: - work._closed = True - work._submitting = False - work._dependencies = () - work._plan = None - work._recovery_token = None - raise - - def _submit_indexed_selection( - self, - *, - plan: TransferPlan | None, - op: TransferOp | str | None, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - selection: IndexSelection | None, - notification: bytes | bytearray | memoryview | None, - cuda_ordering: CudaOrderingMode | str, - stream: object | None, - ready_event: object | None, - recovery_token: object | None, - initial_max_polls: int | None = None, - initial_timeout_ns: int | None = None, - ) -> Work: - """Submit a reusable public or one-shot private indexed selection.""" - - active_plan: TransferPlan | None = None - active_op: TransferOp | None = None - backend_selection: BackendIndexSelection | None = None - cuda_device: object | None = None - submission_token: object | None = None - work: Work | None = None - try: - state_lock = self._state_lock - if state_lock is _NOOP_LOCK: - ( - active_plan, - active_op, - backend_selection, - notification, - cuda_ordering, - cuda_device, - work, - ) = self._reserve_indexed_selection_submission( - plan=plan, - op=op, - local_indices=local_indices, - remote_indices=remote_indices, - selection=selection, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - submission_token = work - self._works.add(work) - active_plan._submitting.add(work) - active_plan._works.add(work) - else: - with state_lock: - ( - active_plan, - active_op, - backend_selection, - notification, - cuda_ordering, - cuda_device, - work, - ) = self._reserve_indexed_selection_submission( - plan=plan, - op=op, - local_indices=local_indices, - remote_indices=remote_indices, - selection=selection, - notification=notification, - cuda_ordering=cuda_ordering, - stream=stream, - ready_event=ready_event, - recovery_token=recovery_token, - ) - submission_token = work - self._works.add(work) - active_plan._submitting.add(work) - active_plan._works.add(work) - - if cuda_ordering is not CudaOrderingMode.CALLER_READY: - _establish_cuda_ordering_prevalidated( - cuda_ordering, - stream, - ready_event, - cuda_device, - ) - if initial_max_polls is None: - result = self._backend_submit_prevalidated( - active_op, - active_plan._handle, - selection=backend_selection, - notification=notification, - # A direct selection submit is the allocation-minimal - # dynamic form used once by framework requests. Bind - # explicitly for native reuse across generations. - reuse_token=None, - adopt_work=work._adopt_backend_work, - ) - if result is not None: - raise BackendFailureError( - "backend prevalidated submit must return None after " - "adopting work" - ) - else: - result = self._backend_submit_prevalidated_and_poll( - active_op, - active_plan._handle, - selection=backend_selection, - notification=notification, - max_polls=initial_max_polls, - timeout_ns=initial_timeout_ns, - adopt_work=work._adopt_backend_work, - ) - current_work = work - work_lock = current_work._lock - if work_lock is None: - if current_work._backend_work is None: - raise BackendFailureError( - "backend prevalidated submit returned without adopting work" - ) - if initial_max_polls is not None: - current_work._ingest_backend_state_unlocked(result) - else: - with work_lock: - if current_work._backend_work is None: - raise BackendFailureError( - "backend prevalidated submit returned without adopting work" - ) - if initial_max_polls is not None: - current_work._ingest_backend_state_unlocked(result) - if state_lock is _NOOP_LOCK: - active_plan._submitting.discard(submission_token) - current_work._submitting = False - else: - with state_lock: - active_plan._submitting.discard(submission_token) - with current_work._lock or _NOOP_LOCK: - current_work._submitting = False - return current_work - except BaseException as error: - state_lock = self._state_lock - if state_lock is _NOOP_LOCK: - adopted = self._reconcile_failed_indexed_submission( - active_plan, work, submission_token, error - ) - else: - with state_lock: - adopted = self._reconcile_failed_indexed_submission( - active_plan, work, submission_token, error - ) - if adopted: - return cast(Work, work) - raise - - def _reconcile_failed_indexed_submission( - self, - active_plan: TransferPlan | None, - work: Work | None, - submission_token: object | None, - error: BaseException, - ) -> bool: - """Repair indexed submission while the caller holds state lock.""" - - if active_plan is not None and work is not None: - active_plan._submitting.discard( - work if submission_token is None else submission_token - ) - if work is None: - return False - work_lock = work._lock - if work_lock is None: - return self._reconcile_failed_indexed_work_unlocked( - active_plan, work, error - ) - with work_lock: - return self._reconcile_failed_indexed_work_unlocked( - active_plan, work, error - ) - - def _reconcile_failed_indexed_work_unlocked( - self, - active_plan: TransferPlan | None, - work: Work, - error: BaseException, - ) -> bool: - if work._backend_work is not None: - work._submission_error = error - if work._state in _TERMINAL_STATES: - work._state = WorkState.FAILED - work._error = error - work._submitting = False - return True - if active_plan is not None: - active_plan._works.discard(work) - self._works.discard(work) - work._closed = True - work._submitting = False - work._dependencies = () - work._plan = None - work._recovery_token = None - return False - - def _submit_bound( - self, - bound: BoundTransfer, - *, - notification: object = _UNSET_NOTIFICATION, - recovery_token: object | None, - ) -> Work: - """Submit a binding without repeating layout or storage validation.""" - - plan: TransferPlan | None = None - selection: IndexSelection | None = None - active_notification: bytes | None = None - reuse_token: object | None = None - work: Work | None = None - try: - state_lock = self._state_lock - with state_lock: - self._check_open() - if not isinstance(bound, BoundTransfer): - raise TypeError("expected a BoundTransfer") - if bound._endpoint is not self: - raise InvalidRegionError( - "bound transfer belongs to another endpoint" - ) - if bound._closed or bound._plan is None: - raise ClosedError("bound transfer is closed") - plan = bound._plan - if plan._closed: - raise ClosedError("transfer plan is closed") - if plan._creation_failed: - raise BackendFailureError( - "transfer plan is retained for cleanup only" - ) - selection = bound._selection - if selection is None: - if notification is not _UNSET_NOTIFICATION: - raise TypeError( - "notification override requires an IndexSelection binding" - ) - active_notification = bound._notification - else: - active_notification = ( - None - if notification is _UNSET_NOTIFICATION - else _validate_attached_notification( - self.capabilities, - self.backend, - notification, - ) - ) - reuse_token = selection._reuse_token - if reuse_token is None: - reuse_token = _ReuseToken() - object.__setattr__(selection, "_reuse_token", reuse_token) - work = Work( - self, - None, - plan=plan, - bound=bound, - dependencies=( - () if active_notification is None else (active_notification,) - ), - recovery_token=recovery_token, - _token=_CONSTRUCTION_TOKEN, - ) - # Use Work itself as the in-progress token. This publishes all - # lifetime reservations before CUDA or provider activity while - # avoiding a second per-submission allocation. - self._works.add(work) - plan._submitting.add(work) - bound._submitting.add(work) - plan._works.add(work) - bound._works.add(work) - - if bound._cuda_ordering is not CudaOrderingMode.CALLER_READY: - _establish_cuda_ordering_prevalidated( - bound._cuda_ordering, - bound._stream, - bound._ready_event, - bound._cuda_device, - ) - - if selection is None: - result = self._backend_submit( - bound._op, - plan._handle, - local_indices=bound._local_indices, - remote_indices=bound._remote_indices, - notification=active_notification, - reuse_token=bound._reuse_token, - adopt_work=work._adopt_backend_work, - ) - else: - result = self._backend_submit_prevalidated( - bound._op, - plan._handle, - selection=selection._backend_selection, - notification=active_notification, - reuse_token=reuse_token, - adopt_work=work._adopt_backend_work, - ) - if result is not None: - raise BackendFailureError( - "backend submit must return None after adopting its work" - ) - current_work = work - work_lock = current_work._lock - if work_lock is None: - if current_work._backend_work is None: - raise BackendFailureError( - "backend submit returned without adopting transfer work" - ) - else: - with work_lock: - if current_work._backend_work is None: - raise BackendFailureError( - "backend submit returned without adopting transfer work" - ) - if state_lock is _NOOP_LOCK: - plan._submitting.discard(current_work) - bound._submitting.discard(current_work) - current_work._submitting = False - else: - with state_lock: - plan._submitting.discard(current_work) - bound._submitting.discard(current_work) - with current_work._lock or _NOOP_LOCK: - current_work._submitting = False - return current_work - except BaseException as error: - with self._state_lock: - if plan is not None and work is not None: - plan._submitting.discard(work) - bound._submitting.discard(work) - adopted = False - if work is not None: - with work._lock or _NOOP_LOCK: - adopted = work._backend_work is not None - if adopted: - # Preserve the generic submission recovery contract: once - # adopted, an interruption is reported by the retained Work. - work = cast(Work, work) - with work._lock or _NOOP_LOCK: - work._submission_error = error - work._submitting = False - return work - if work is not None: - if plan is not None: - plan._works.discard(work) - bound._works.discard(work) - self._works.discard(work) - with work._lock or _NOOP_LOCK: - work._closed = True - work._submitting = False - work._dependencies = () - work._plan = None - work._bound = None - work._recovery_token = None - raise - - def send_notification( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - *, - recovery_token: object | None = None, - ) -> Work: - if self._thread_mode is ThreadMode.MULTIPLE: - return self._send_notification_multiple( - peer, payload, recovery_token=recovery_token - ) - - state_lock = self._state_lock - if state_lock is _NOOP_LOCK: - return self._send_notification_locked( - peer, payload, recovery_token=recovery_token - ) - with state_lock: - return self._send_notification_locked( - peer, payload, recovery_token=recovery_token - ) - - def _send_notification_locked( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - *, - recovery_token: object | None, - ) -> Work: - """Send while the selected non-MULTIPLE lifecycle boundary is held.""" - - self._check_open() - self._validate_peer(peer) - if not self.capabilities.standalone_notifications: - raise UnsupportedError( - f"backend {self.backend!r} does not support notifications" - ) - payload = _as_bytes(payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES) - work = Work( - self, - None, - dependencies=(peer, payload), - recovery_token=recovery_token, - _token=_CONSTRUCTION_TOKEN, - ) - try: - self._works.add(work) - result = self._backend_send_notification( - peer._handle, - payload, - adopt_work=work._adopt_backend_work, - ) - if result is not None: - raise BackendFailureError( - "backend notification send must return None after adopting its work" - ) - lock = work._lock - if lock is None: - if work._backend_work is None: - raise BackendFailureError( - "backend notification send returned without adopting work" - ) - else: - with lock: - if work._backend_work is None: - raise BackendFailureError( - "backend notification send returned without adopting work" - ) - work._submitting = False - return work - except BaseException as error: - with work._lock or _NOOP_LOCK: - adopted = work._backend_work is not None - if adopted: - work._submission_error = error - work._submitting = False - return work - work._closed = True - work._submitting = False - work._dependencies = () - work._recovery_token = None - self._works.discard(work) - raise - - def _send_notification_multiple( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - *, - recovery_token: object | None, - ) -> Work: - """Publish Work ownership, then preserve provider MULTIPLE concurrency.""" - - work: Work | None = None - try: - with self._state_lock: - self._check_open() - self._validate_peer(peer) - if not self._capabilities.standalone_notifications: - raise UnsupportedError( - f"backend {self._backend_name!r} does not support notifications" - ) - payload = _as_bytes( - payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES - ) - work = Work( - self, - None, - dependencies=(peer, payload), - recovery_token=recovery_token, - _token=_CONSTRUCTION_TOKEN, - ) - # Work itself is the allocation already required by this API and - # becomes the close-visible lifetime reservation before native - # entry. No second per-send token is needed. - self._works.add(work) - - result = self._backend_send_notification( - peer._handle, - payload, - adopt_work=work._adopt_backend_work, - ) - if result is not None: - raise BackendFailureError( - "backend notification send must return None after adopting its work" - ) - with work._lock or _NOOP_LOCK: - if work._backend_work is None: - raise BackendFailureError( - "backend notification send returned without adopting work" - ) - with self._state_lock: - with work._lock or _NOOP_LOCK: - work._submitting = False - return work - except BaseException as error: - with self._state_lock: - if work is None: - raise - with work._lock or _NOOP_LOCK: - if work._backend_work is not None: - work._submission_error = error - work._submitting = False - return work - work._closed = True - work._submitting = False - work._dependencies = () - work._recovery_token = None - self._works.discard(work) - raise - - def send_notification_sync( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - ) -> None: - """Send without allocating a :class:`Work`. - - Return means exactly ``standalone_notification_completion``. In - particular, NIXL's ``LOCAL_ACCEPTED`` does not mean the peer observed - the payload. Once backend entry begins, an exception can be - outcome-ambiguous; protocols that need certainty must use sequence - numbers and acknowledgement, or retain :meth:`send_notification` and - its inspectable Work ownership. - """ - - if self._thread_mode is ThreadMode.MULTIPLE: - self._send_notification_sync_multiple(peer, payload) - return - - state_lock = self._state_lock - if state_lock is _NOOP_LOCK: - self._send_notification_sync_locked(peer, payload) - return - with state_lock: - self._send_notification_sync_locked(peer, payload) - - def _send_notification_sync_locked( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - ) -> None: - """Validate and send while the caller supplies lifecycle exclusion.""" - - self._check_open() - self._validate_peer(peer) - if not self._capabilities.synchronous_notification_send: - raise UnsupportedError( - f"backend {self._backend_name!r} does not support synchronous " - "notification send" - ) - payload = _as_bytes(payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES) - send = self._send_notification_sync_dispatch - if send is None: - raise BackendFailureError( - "backend advertised synchronous notification send without a " - "dispatch method" - ) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = send(peer._handle, payload) - else: - result = send(peer._handle, payload) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._backend_name!r} failed to send notification " - "synchronously; delivery outcome may be ambiguous" - ) from error - if result is not None: - raise BackendFailureError( - "backend synchronous notification send must return None; " - "delivery outcome may be ambiguous" - ) - - def _send_notification_sync_multiple( - self, - peer: Peer, - payload: bytes | bytearray | memoryview, - ) -> None: - """Reserve peer lifetime briefly, then preserve native RW concurrency.""" - - reserved_peer: Peer | None = None - reservation: _SyncNotificationReservation | None = None - try: - with self._state_lock: - self._check_open() - self._validate_peer(peer) - if not self._capabilities.synchronous_notification_send: - raise UnsupportedError( - f"backend {self._backend_name!r} does not support synchronous " - "notification send" - ) - payload = _as_bytes( - payload, "payload", max_nbytes=_MAX_NOTIFICATION_BYTES - ) - send = self._send_notification_sync_dispatch - if send is None: - raise BackendFailureError( - "backend advertised synchronous notification send without a " - "dispatch method" - ) - peer_handle = peer._handle - pool = peer._sync_notification_reservation_pool - reservation = pool.pop() if pool else _SyncNotificationReservation() - reserved_peer = peer - reservation.active = True - peer._active_sync_notification_sends.add(reservation.reference) - - try: - result = send(peer_handle, payload) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self._backend_name!r} failed to send notification " - "synchronously; delivery outcome may be ambiguous" - ) from error - if result is not None: - raise BackendFailureError( - "backend synchronous notification send must return None; " - "delivery outcome may be ambiguous" - ) - finally: - if reserved_peer is not None and reservation is not None: - # Publish native exit before taking an interruptible cleanup - # lock. A retained exception traceback may keep this frame and - # token alive, so liveness of the weakref alone is insufficient. - reservation.active = False - with self._state_lock: - reserved_peer._active_sync_notification_sends.discard( - reservation.reference - ) - if not reserved_peer._closed: - reserved_peer._sync_notification_reservation_pool.append( - reservation - ) - - def poll_notifications(self, max_items: int | None = None) -> list[Notification]: - if self._thread_mode is ThreadMode.MULTIPLE: - notifications = self._poll_backend_notifications_multiple( - max_items, grouped=False - ) - else: - state_lock = self._state_lock - with state_lock: - self._check_open() - if not self.capabilities.notifications: - raise UnsupportedError( - f"backend {self.backend!r} does not support notifications" - ) - if max_items is not None: - if isinstance(max_items, bool) or not isinstance(max_items, int): - raise TypeError( - "max_items must be a non-negative integer or None" - ) - if max_items < 0: - raise ValueError("max_items must be non-negative") - if max_items == 0: - return [] - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - notifications = self._backend.poll_notifications(max_items) - else: - notifications = self._backend.poll_notifications(max_items) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to poll notifications" - ) from error - if not isinstance(notifications, list) or not all( - isinstance(item, BackendNotification) for item in notifications - ): - raise BackendFailureError( - "backend notifications must be a list of BackendNotification" - ) - if max_items is not None and len(notifications) > max_items: - raise BackendFailureError( - "backend returned more notifications than max_items" - ) - return [ - Notification( - payload=notification.payload, - source_endpoint_id=notification.source_endpoint_id, - source_incarnation=notification.source_incarnation, - ) - for notification in notifications - ] - - def poll_notification_batches( - self, max_items: int | None = None - ) -> list[NotificationBatch]: - """Destructively receive source-grouped notifications. - - ``max_items`` counts payloads rather than returned batches. The v0 - receive contract remains at-most-once: there is no acknowledgement or - rollback token. Per-source payload order is preserved, but providers - need not define a total order across different sources and may return - repeated batches for one source. - """ - - if self._thread_mode is ThreadMode.MULTIPLE: - batches = self._poll_backend_notifications_multiple(max_items, grouped=True) - else: - state_lock = self._state_lock - with state_lock: - self._check_open() - if not self.capabilities.notifications: - raise UnsupportedError( - f"backend {self.backend!r} does not support notifications" - ) - if max_items is not None: - if isinstance(max_items, bool) or not isinstance(max_items, int): - raise TypeError( - "max_items must be a non-negative integer or None" - ) - if max_items < 0: - raise ValueError("max_items must be non-negative") - if max_items == 0: - return [] - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - batches = self._backend.poll_notification_batches(max_items) - else: - batches = self._backend.poll_notification_batches(max_items) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to poll notification batches" - ) from error - if not isinstance(batches, list) or not all( - type(item) is BackendNotificationBatch for item in batches - ): - raise BackendFailureError( - "backend notification batches must be a list of " - "BackendNotificationBatch" - ) - if max_items is not None: - payload_count = sum(len(batch.payloads) for batch in batches) - if payload_count > max_items: - raise BackendFailureError( - "backend returned more notification payloads than max_items" - ) - return [_notification_batch_from_backend(batch) for batch in batches] - - def _poll_backend_notifications_multiple( - self, - max_items: int | None, - *, - grouped: bool, - ) -> object: - """Run one destructive receive with narrow, close-visible admission.""" - - reservation: _BackendCallReservation | None = None - # Serialize only destructive consumers. Acquire this lock before call - # admission so queued consumers need no reservation and can safely - # discover a concurrently completed close before backend entry. - with self._notification_consumer_lock: - try: - with self._state_lock: - self._check_open() - if not self._capabilities.notifications: - raise UnsupportedError( - f"backend {self._backend_name!r} does not support " - "notifications" - ) - if max_items is not None: - if isinstance(max_items, bool) or not isinstance( - max_items, int - ): - raise TypeError( - "max_items must be a non-negative integer or None" - ) - if max_items < 0: - raise ValueError("max_items must be non-negative") - if max_items == 0: - return [] - poll = ( - self._backend.poll_notification_batches - if grouped - else self._backend.poll_notifications - ) - reservation = self._reserve_multiple_backend_call_locked() - try: - return poll(max_items) - except TransferError: - raise - except Exception as error: - action = ( - "poll notification batches" if grouped else "poll notifications" - ) - raise BackendFailureError( - f"backend {self._backend_name!r} failed to {action}" - ) from error - finally: - if reservation is not None: - self._finish_multiple_backend_call(reservation) - - def progress(self) -> None: - if self._thread_mode is ThreadMode.MULTIPLE: - self._progress_multiple() - return - - state_lock = self._state_lock - if state_lock is _NOOP_LOCK: - self._progress_locked() - return - with state_lock: - self._progress_locked() - - def _progress_locked(self) -> None: - self._check_open() - if not self.capabilities.manual_progress: - raise UnsupportedError( - f"backend {self.backend!r} does not support manual progress" - ) - self._backend_call("progress backend", self._backend.progress) - - def _progress_multiple(self) -> None: - reservation: _BackendCallReservation | None = None - try: - with self._state_lock: - self._check_open() - if not self._capabilities.manual_progress: - raise UnsupportedError( - f"backend {self._backend_name!r} does not support manual progress" - ) - progress = self._backend.progress - reservation = self._reserve_multiple_backend_call_locked() - self._backend_call("progress backend", progress) - finally: - if reservation is not None: - self._finish_multiple_backend_call(reservation) - - @_state_locked - def close(self) -> None: - if self._creation_failed: - self._close_failed_endpoint_backend() - return - if self._closed: - return - if not self._closing: - if self._multiple_backend_calls_active_locked(): - raise BusyError( - "cannot close endpoint while backend calls are in progress" - ) - self._retry_orphaned_plans() - active = [work for work in self._works if not work.test()] - if active: - raise BusyError( - f"cannot close endpoint with {len(active)} active work item(s)" - ) - counts = { - "work": len(self._works), - "transfer slot": len(self._transfer_slots), - "bound transfer": len(self._bound_transfers), - "plan": len(self._plans), - "peer": len(self._peers), - "registration": len(self._registrations), - } - if any(counts.values()): - detail = ", ".join( - f"{name}={count}" for name, count in counts.items() if count - ) - raise BusyError( - f"close child resources before endpoint close ({detail})" - ) - # From this point only close() may be retried. This makes a signal - # after successful provider teardown converge without calling - # progress on an already closed provider. - self._closing = True - self._backend_call( - "close backend", - lambda: self._backend_receipt.close(self._backend.close), - ) - self._detach_backend_finalizer() - self._closed = True - active_calls = self._active_multiple_backend_calls - if active_calls is not None: - active_calls.clear() - reservation_pool = self._multiple_backend_call_reservation_pool - if reservation_pool is not None: - reservation_pool.clear() - - def _close_failed_endpoint_backend(self) -> None: - """Retry backend-only cleanup after Endpoint initialization failed.""" - - if self._closed: - self._creation_failed = False - self._detach_backend_finalizer() - return - self._backend_call( - "close backend after failed endpoint initialization", - lambda: self._backend_receipt.close(self._backend.close), - ) - self._detach_backend_finalizer() - self._closed = True - self._creation_failed = False - - def _cleanup_adopted_backend_after_creation_error( - self, original_error: BaseException - ) -> None: - implementation = self._backend_receipt.backend - if implementation is None: - self._detach_backend_finalizer() - return - self._backend = implementation - self._creation_failed = True - try: - self._close_failed_endpoint_backend() - except BaseException as cleanup_error: - raise _endpoint_creation_cleanup_failure( - self, original_error, cleanup_error - ) from cleanup_error - - def _detach_backend_finalizer(self) -> None: - self._backend_finalizer.detach() - - def _submit( - self, - op: TransferOp, - plan: TransferPlan | None, - *, - local: Sequence[Region] | None, - remote: Sequence[Region] | None, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | bytearray | memoryview | None, - cuda_ordering: CudaOrderingMode | str, - stream: object | None, - ready_event: object | None, - recovery_token: object | None, - ) -> Work: - owned_plan = False - plan_to_release: TransferPlan | None = None - reserved_plan: TransferPlan | None = None - submission_token: object | None = None - work: Work | None = None - try: - state_lock = self._state_lock - with state_lock: - self._check_open() - supported = ( - self.capabilities.read - if op is TransferOp.READ - else self.capabilities.write - ) - if not supported: - raise UnsupportedError( - f"backend {self.backend!r} does not support {op.value}" - ) - cuda_ordering = _coerce_enum( - cuda_ordering, CudaOrderingMode, "cuda_ordering" - ) - _validate_cuda_ordering_arguments( - self.capabilities, - self.backend, - cuda_ordering, - stream, - ready_event, - ) - if plan is None: - if local is None or remote is None: - raise TypeError( - "one-shot transfer requires both local and remote regions" - ) - plan = self.prepare(local=local, remote=remote, indexed=False) - plan_to_release = plan - owned_plan = True - elif local is not None or remote is not None: - raise TypeError("pass a plan or one-shot regions, not both") - active_plan = plan - self._validate_plan_locked(active_plan) - local_indices, remote_indices = _validate_plan_indices( - active_plan, local_indices, remote_indices, op - ) - notification = _validate_attached_notification( - self.capabilities, self.backend, notification - ) - cuda_device = _validate_cuda_ordering_selection( - active_plan, - local_indices, - cuda_ordering, - stream, - ready_event, - ) - dependencies = tuple( - value - for value in ( - local_indices, - remote_indices, - stream, - ready_event, - notification, - ) - if value is not None - ) - # Publish the complete Core lifetime graph before entering the - # provider. The provider adopts its request into this wrapper - # before it is allowed to launch transport work. - work = Work( - self, - None, - plan=active_plan, - dependencies=dependencies, - owned_plan=owned_plan, - recovery_token=recovery_token, - _token=_CONSTRUCTION_TOKEN, - ) - submission_token = work - reserved_plan = active_plan - self._works.add(work) - active_plan._submitting.add(submission_token) - active_plan._works.add(work) - _establish_cuda_ordering_prevalidated( - cuda_ordering, - stream, - ready_event, - cuda_device, - ) - - reuse_token = ( - active_plan._reuse_tokens[op] - if not owned_plan - and local_indices is None - and remote_indices is None - and notification is None - and cuda_ordering is CudaOrderingMode.CALLER_READY - else None - ) - result = self._backend_submit( - op, - active_plan._handle, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - reuse_token=reuse_token, - adopt_work=work._adopt_backend_work, - ) - if result is not None: - raise BackendFailureError( - "backend submit must return None after adopting its work" - ) - current_work = work - work_lock = current_work._lock - if work_lock is None: - if current_work._backend_work is None: - raise BackendFailureError( - "backend submit returned without adopting transfer work" - ) - else: - with work_lock: - if current_work._backend_work is None: - raise BackendFailureError( - "backend submit returned without adopting transfer work" - ) - if state_lock is _NOOP_LOCK: - reserved_plan._submitting.discard(submission_token) - current_work._submitting = False - else: - with state_lock: - reserved_plan._submitting.discard(submission_token) - with current_work._lock or _NOOP_LOCK: - current_work._submitting = False - return current_work - except BaseException as error: - with self._state_lock: - if reserved_plan is not None and submission_token is not None: - reserved_plan._submitting.discard(submission_token) - adopted = False - if work is not None: - with work._lock or _NOOP_LOCK: - adopted = work._backend_work is not None - if adopted: - # A provider must not raise after adoption, but retaining - # the Work is the only safe response if it does: native - # access may already be in flight. Surface the protocol - # violation only after the backend request is terminal. - work = cast(Work, work) - with work._lock or _NOOP_LOCK: - work._submission_error = error - work._submitting = False - return work - if work is not None: - if work._plan is not None: - work._plan._works.discard(work) - self._works.discard(work) - with work._lock or _NOOP_LOCK: - work._closed = True - work._submitting = False - work._dependencies = () - work._plan = None - work._recovery_token = None - if plan_to_release is not None: - try: - plan_to_release.close() - except BaseException: - self._orphaned_plans.add(plan_to_release) - raise - - def _validate_regions( - self, - local: tuple[Region, ...], - remote: tuple[Region, ...], - *, - indexed: bool, - ) -> None: - if not local or not remote: - raise InvalidRegionError("local and remote regions must be non-empty") - for region in local: - if not isinstance(region, Region) or region.is_remote: - raise InvalidRegionError( - "local must contain local Registration regions" - ) - if not isinstance(region._origin, Registration): - raise InvalidRegionError("local region has an invalid owner") - self._validate_registration(region._origin) - peers: set[Peer] = set() - for region in remote: - if not isinstance(region, Region) or not region.is_remote: - raise InvalidRegionError("remote must contain Peer regions") - if not isinstance(region._origin, Peer): - raise InvalidRegionError("remote region has an invalid owner") - self._validate_peer(region._origin) - peers.add(region._origin) - if len(peers) != 1: - raise InvalidRegionError("one transfer plan may target exactly one peer") - if not indexed: - local_size = sum(region.total_nbytes for region in local) - remote_size = sum(region.total_nbytes for region in remote) - if local_size != remote_size: - raise InvalidRegionError( - f"local and remote byte totals differ: " - f"{local_size} != {remote_size}" - ) - - def _validate_registration(self, registration: Registration) -> None: - if not isinstance(registration, Registration): - raise TypeError("expected a Registration") - if registration._endpoint is not self: - raise InvalidRegionError("registration belongs to another endpoint") - registration._check_open() - _validate_memory_stable(registration._memory) - - def _validate_peer(self, peer: Peer) -> None: - if not isinstance(peer, Peer): - raise TypeError("expected a Peer") - if peer._endpoint is not self: - raise InvalidRegionError("peer belongs to another endpoint") - peer._check_open() - - def _validate_plan(self, plan: TransferPlan) -> None: - self._check_open() - self._validate_plan_locked(plan) - - def _validate_plan_locked(self, plan: TransferPlan) -> None: - """Validate a plan after this endpoint has entered its state lock.""" - - if not isinstance(plan, TransferPlan): - raise TypeError("expected a TransferPlan") - if plan._endpoint is not self: - raise InvalidRegionError("transfer plan belongs to another endpoint") - if plan._closed: - raise ClosedError("transfer plan is closed") - if plan._creation_failed: - raise BackendFailureError("transfer plan is retained for cleanup only") - for origin in plan._origins: - if isinstance(origin, Registration): - if origin._endpoint is not self: - raise InvalidRegionError("registration belongs to another endpoint") - if origin._closed: - raise ClosedError("registration is closed") - if origin._creation_failed: - raise BackendFailureError( - "registration is retained for cleanup only" - ) - _validate_memory_stable(origin._memory) - - @_state_locked - def _close_registration(self, registration: Registration) -> None: - if registration._closed: - self._finish_registration_close(registration) - return - self._check_open() - self._retry_orphaned_plans() - if any( - not plan._closed and registration in plan._origins for plan in self._plans - ): - raise BusyError( - f"registration {registration.name!r} is used by an open plan" - ) - result = self._backend_call( - "deregister memory", - lambda: self._backend.deregister(registration._handle), - ) - if result is not None: - raise BackendFailureError("backend deregister must return None") - registration._closed = True - self._finish_registration_close(registration) - - def _finish_registration_close(self, registration: Registration) -> None: - """Idempotently finish Core-only registration teardown.""" - - registration_id = registration.registration_id - name = registration.name - if self._registration_names.get(name) is registration: - self._registration_names.pop(name, None) - registration._handle = None - registration._memory = _released_memory(registration._memory) - # Authoritative recovery registry removal is the Core teardown commit. - if self._registrations.get(registration_id) is registration: - self._registrations.pop(registration_id, None) - - @_state_locked - def _close_peer(self, peer: Peer) -> None: - if peer._closed: - self._finish_peer_close(peer) - return - self._check_open() - self._retry_orphaned_plans() - active_sends = peer._active_sync_notification_sends - for reference in tuple(active_sends): - reservation = reference() - if reservation is None or not reservation.active: - active_sends.discard(reference) - if active_sends: - raise BusyError( - f"peer {peer.name!r} has synchronous notification sends in progress" - ) - if any(not plan._closed and peer in plan._origins for plan in self._plans): - raise BusyError(f"peer {peer.name!r} is used by an open plan") - if any(peer in work._dependencies for work in self._works): - raise BusyError(f"peer {peer.name!r} is used by an open work item") - result = self._backend_call( - "release peer", lambda: self._backend.release_peer(peer._handle) - ) - if result is not None: - raise BackendFailureError("backend release_peer must return None") - peer._closed = True - self._finish_peer_close(peer) - - def _finish_peer_close(self, peer: Peer) -> None: - """Idempotently finish Core-only peer teardown.""" - - if self._peers.get(peer.endpoint_id) is peer: - self._peers.pop(peer.endpoint_id, None) - peer._handle = None - peer._registrations = {} - peer._active_sync_notification_sends.clear() - peer._sync_notification_reservation_pool.clear() - - @_state_locked - def _close_plan(self, plan: TransferPlan) -> None: - if plan._closed: - self._finish_plan_close(plan) - return - self._check_open() - if plan._submitting: - raise BusyError("transfer plan has submissions in progress") - if plan._works: - raise BusyError("transfer plan has open work items") - if plan._bindings: - raise BusyError("transfer plan has open bound transfers") - result = self._backend_call( - "release transfer plan", - lambda: self._backend.release_plan(plan._handle), - ) - if result is not None: - raise BackendFailureError("backend release_plan must return None") - plan._closed = True - self._finish_plan_close(plan) - - def _finish_plan_close(self, plan: TransferPlan) -> None: - """Idempotently finish Core-only transfer-plan teardown.""" - - self._plans.discard(plan) - self._orphaned_plans.discard(plan) - plan._handle = None - plan._local = () - plan._remote = () - plan._origins = frozenset() - plan._reuse_tokens = {} - - @_state_locked - def _close_bound_transfer(self, bound: BoundTransfer) -> None: - if bound._closed: - # Closure is the commit marker, not an early-return condition. An - # interrupt after it may leave a parent edge or retained field; - # retrying completes the idempotent Core-only epilogue. - self._finish_bound_transfer_close(bound) - return - self._check_open() - if bound._endpoint is not self: - raise InvalidRegionError("bound transfer belongs to another endpoint") - plan = bound._check_open() - if bound._slots: - raise BusyError("bound transfer has open persistent request slots") - if bound._works: - raise BusyError("bound transfer has open work items") - if bound._submitting: - raise BusyError("bound transfer has submissions in progress") - # Unlink blocking graph edges before publishing CLOSED, while retaining - # `_plan` so a retry after that commit can repair either edge. - plan._bindings.discard(bound) - self._bound_transfers.discard(bound) - bound._closed = True - self._finish_bound_transfer_close(bound) - - def _finish_bound_transfer_close(self, bound: BoundTransfer) -> None: - """Idempotently finish Core-only bound-transfer unlink.""" - - plan = bound._plan - if plan is not None: - plan._bindings.discard(bound) - self._bound_transfers.discard(bound) - bound._plan = None - bound._local_indices = None - bound._remote_indices = None - bound._notification = None - bound._cuda_device = None - bound._stream = None - bound._ready_event = None - bound._reuse_token = None - bound._selection = None - - @_state_locked - def _close_transfer_slot(self, slot: TransferSlot) -> None: - """Release one idle persistent request and unlink its lifetime graph.""" - - if not isinstance(slot, TransferSlot): - raise TypeError("expected a TransferSlot") - if slot._endpoint is not self: - raise InvalidRegionError("transfer slot belongs to another endpoint") - if slot._closed: - # Closure is the commit marker, not an early-return condition. An - # interrupt after it may still leave one parent edge or retained - # callback, so every retry finishes the idempotent Core epilogue. - self._finish_transfer_slot_close(slot) - return - self._check_open() - bound = slot._bound - if bound is None or bound._closed: - raise ClosedError("transfer slot's bound transfer is closed") - if slot._active: - raise BusyError( - "transfer slot owns an unrecycled generation; recycle it before close" - ) - close = slot._backend_close - if close is None: - if slot._backend_slot is not None: - raise BackendFailureError( - "transfer slot has no backend close operation" - ) - else: - self._backend_request_slot_close(close) - - # Closure commits provider release while both recovery edges still - # exist. The retry path can therefore remove the bound blocker first - # and the endpoint's public recovery root last. - slot._closed = True - self._finish_transfer_slot_close(slot) - - def _finish_transfer_slot_close(self, slot: TransferSlot) -> None: - """Idempotently finish Core-only unlink after provider release.""" - - bound = slot._bound - if bound is not None: - bound._slots.discard(slot) - self._transfer_slots.discard(slot) - slot._bound = None - slot._backend_slot = None - slot._backend_execution_dispatch = None - slot._backend_start = None - slot._backend_start_and_poll = None - if isinstance(slot, RequestLocalNotificationSlot): - slot._backend_start_with_notification = None - slot._backend_start_and_poll_with_notification = None - slot._notification_fast_execute = None - slot._backend_poll_bounded = None - slot._backend_poll_state = None - slot._backend_cancel = None - slot._backend_recycle = None - slot._backend_close = None - slot._ready_event_synchronize = None - slot._fast_execute = None - slot._backend_generation_active = False - slot._error = None - slot._start_error = None - - def _retry_orphaned_plans(self) -> None: - for plan in tuple(self._orphaned_plans): - plan.close() - - def _check_open(self) -> None: - self._check_thread() - if self._creation_failed: - raise BackendFailureError( - f"endpoint {self.name!r} is retained for creation cleanup only; " - "retry close()" - ) - if self._closed: - raise ClosedError(f"endpoint {self.name!r} is closed") - if self._closing: - raise ClosedError( - f"endpoint {self.name!r} is closing; only close() may be retried" - ) - - def _reserve_multiple_backend_call_locked(self) -> _BackendCallReservation: - active_calls = self._active_multiple_backend_calls - pool = self._multiple_backend_call_reservation_pool - assert active_calls is not None and pool is not None - reservation = pool.pop() if pool else _BackendCallReservation() - try: - reservation.active = True - active_calls.add(reservation.reference) - except BaseException: - reservation.active = False - active_calls.discard(reservation.reference) - pool.append(reservation) - raise - return reservation - - def _finish_multiple_backend_call( - self, reservation: _BackendCallReservation - ) -> None: - # Publish native exit before interruptible cleanup. A traceback can keep - # this frame alive; close therefore also checks the explicit active bit. - reservation.active = False - with self._state_lock: - active_calls = self._active_multiple_backend_calls - pool = self._multiple_backend_call_reservation_pool - assert active_calls is not None and pool is not None - active_calls.discard(reservation.reference) - if not self._closing and not self._closed: - pool.append(reservation) - - def _multiple_backend_calls_active_locked(self) -> bool: - active_calls = self._active_multiple_backend_calls - if not active_calls: - return False - for reference in tuple(active_calls): - reservation = reference() - if reservation is None or not reservation.active: - active_calls.discard(reference) - return bool(active_calls) - - def _backend_call(self, action: str, function): # type: ignore[no-untyped-def] - self._check_thread() - try: - if self.thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return function() - return function() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to {action}" - ) from error - - def _backend_submit( - self, - op: TransferOp, - plan_handle: object, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | None, - reuse_token: object | None, - adopt_work: Callable[[BackendWork], None], - ) -> object | None: - """Call the submission SPI without allocating an extra lambda.""" - - # Both callers validate endpoint/thread state before reaching this - # same-stack helper; avoid repeating the thread-ident lookup here. - submit = cast(Callable[..., object | None], self._backend.submit) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return submit( - op, - plan_handle, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - reuse_token=reuse_token, - adopt_work=adopt_work, - ) - return submit( - op, - plan_handle, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - reuse_token=reuse_token, - adopt_work=adopt_work, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to submit {op.value}" - ) from error - - def _backend_send_notification( - self, - peer_handle: object, - payload: bytes, - *, - adopt_work: Callable[[BackendWork], None], - ) -> object | None: - """Call standalone notification send without per-call closures.""" - - # send_notification() has already validated endpoint and thread state. - send = cast(Callable[..., object | None], self._backend.send_notification) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return send(peer_handle, payload, adopt_work=adopt_work) - return send(peer_handle, payload, adopt_work=adopt_work) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to send notification" - ) from error - - def _backend_submit_prevalidated( - self, - op: TransferOp, - plan_handle: object, - *, - selection: BackendIndexSelection, - notification: bytes | None, - reuse_token: object | None, - adopt_work: Callable[[BackendWork], None], - ) -> object | None: - """Call the authenticated indexed-selection SPI without a lambda.""" - - submit = cast(Callable[..., object | None], self._backend.submit_prevalidated) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return submit( - op, - plan_handle, - selection=selection, - notification=notification, - reuse_token=reuse_token, - adopt_work=adopt_work, - ) - return submit( - op, - plan_handle, - selection=selection, - notification=notification, - reuse_token=reuse_token, - adopt_work=adopt_work, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to submit prevalidated {op.value}" - ) from error - - def _backend_submit_prevalidated_and_poll( - self, - op: TransferOp, - plan_handle: object, - *, - selection: BackendIndexSelection, - notification: bytes | None, - max_polls: int, - timeout_ns: int | None, - adopt_work: Callable[[BackendWork], None], - ) -> WorkState: - """Call the optional fused one-shot selection SPI without a lambda.""" - - submit = cast( - Callable[..., WorkState], - self._backend.submit_prevalidated_and_poll, - ) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return submit( - op, - plan_handle, - selection=selection, - notification=notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - adopt_work=adopt_work, - ) - return submit( - op, - plan_handle, - selection=selection, - notification=notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - adopt_work=adopt_work, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to submit and poll " - f"prevalidated {op.value}" - ) from error - - def _backend_create_request_slot( - self, - op: TransferOp, - plan_handle: object, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | None, - adopt_slot: Callable[..., None], - ) -> object | None: - """Call the optional persistent-request SPI without a lambda.""" - - create = cast(Callable[..., object | None], self._backend.create_request_slot) - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return create( - op, - plan_handle, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - adopt_slot=adopt_slot, - ) - return create( - op, - plan_handle, - local_indices=local_indices, - remote_indices=remote_indices, - notification=notification, - adopt_slot=adopt_slot, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to create a persistent " - f"{op.value} request slot" - ) from error - - def _backend_request_slot_start(self, start: Callable[[], None]) -> None: - """Start a cached backend slot operation on the lowest-overhead path.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = start() - else: - result = start() - if result is not None: - raise BackendFailureError("backend request-slot start must return None") - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to start a persistent request" - ) from error - - def _backend_request_slot_start_with_notification( - self, - start: Callable[[bytes | None], None], - notification: bytes | None, - ) -> None: - """Start a cached request with an explicit notification override.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = start(notification) - else: - result = start(notification) - if result is not None: - raise BackendFailureError( - "backend request-slot start_with_notification must return None" - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to start a persistent request " - "with a notification override" - ) from error - - def _backend_request_slot_start_and_poll( - self, - start_and_poll: Callable[..., WorkState | None], - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - """Enter the optional bounded persistent-request fused operation.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return start_and_poll( - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - return start_and_poll(max_polls=max_polls, timeout_ns=timeout_ns) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed during fused persistent " - "request start/poll" - ) from error - - def _backend_request_slot_start_and_poll_with_notification( - self, - start_and_poll: Callable[..., WorkState | None], - notification: bytes | None, - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - """Enter fused post/poll with a request-local notification override.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return start_and_poll( - notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - return start_and_poll( - notification, - max_polls=max_polls, - timeout_ns=timeout_ns, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed during fused persistent " - "request notification start/poll" - ) from error - - def _backend_request_slot_poll_bounded( - self, - poll_bounded: Callable[..., WorkState | None], - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState | None: - """Enter the optional bounded persistent-request polling operation.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - return poll_bounded( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed during bounded persistent " - "request polling" - ) from error - - def _backend_request_slot_state( - self, poll_state: Callable[[], WorkState | None] - ) -> WorkState | None: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return poll_state() - return poll_state() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to query persistent request state" - ) from error - - def _backend_request_slot_error( - self, backend_slot: BackendRequestSlot - ) -> BaseException | None: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return backend_slot.error - return backend_slot.error - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to query persistent request error" - ) from error - - def _backend_request_slot_completion_event( - self, backend_slot: BackendRequestSlot - ) -> object | None: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return backend_slot.completion_event - return backend_slot.completion_event - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to query persistent request " - "completion event" - ) from error - - def _backend_request_slot_cancel(self, cancel: Callable[[], bool]) -> bool: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return cancel() - return cancel() - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to cancel a persistent request" - ) from error - - def _backend_request_slot_recycle(self, recycle: Callable[[], None]) -> None: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = recycle() - else: - result = recycle() - if result is not None: - raise BackendFailureError( - "backend request-slot recycle must return None" - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to recycle a persistent request" - ) from error - - def _backend_request_slot_close(self, close: Callable[[], None]) -> None: - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = close() - else: - result = close() - if result is not None: - raise BackendFailureError("backend request-slot close must return None") - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to release a persistent request" - ) from error - - def _backend_work_state(self, backend_work: BackendWork) -> WorkState: - """Read a provider Work state without a lambda or duplicate thread check.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return backend_work.state - return backend_work.state - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to query work state" - ) from error - - def _backend_work_poll_bounded( - self, - backend_work: BackendWork, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState: - """Call a fused provider Work poll without a lambda.""" - - poll = backend_work.poll_bounded - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - return poll( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - return poll( - max_polls=max_polls, - timeout_ns=timeout_ns, - timeout_check_interval=timeout_check_interval, - ) - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to poll work" - ) from error - - def _backend_work_close(self, backend_work: BackendWork) -> None: - """Close provider Work without a bound-method/lambda dispatch wrapper.""" - - try: - if self._thread_mode is ThreadMode.SERIALIZED: - with self._backend_lock: - result = backend_work.close() - else: - result = backend_work.close() - if result is not None: - raise BackendFailureError("backend work close must return None") - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - f"backend {self.backend!r} failed to release work" - ) from error - - def _check_thread(self) -> None: - if ( - self._thread_mode is ThreadMode.SINGLE - and threading.get_ident() != self._owner_thread - ): - raise UnsupportedError( - "single-thread endpoint used from a different thread" - ) - - def __enter__(self) -> Self: - self._check_open() - return self - - def __exit__(self, *args: object) -> None: - self.close() - - -def _registered_memory( - value: object, - registration_id: str, - name: str, -) -> RegisteredMemory: - if isinstance(value, RawSpan): - return RegisteredMemory( - registration_id=registration_id, - name=name, - owner=_LifetimeAnchor(value.owner), - address=value.address, - nbytes=value.nbytes, - device=value.device, - memory_type=value.memory_type, - ) - if isinstance(value, torch.Tensor): - if value.layout is not torch.strided: - raise InvalidRegionError("only strided tensor layouts can be registered") - if not value.is_contiguous(): - raise InvalidRegionError( - "tensor registration requires a contiguous tensor; " - "no implicit copy is performed" - ) - nbytes = value.numel() * value.element_size() - address = value.data_ptr() - device = str(value.device) - memory_type = _tensor_memory_type(value) - storage = value.untyped_storage() - anchor = _LifetimeAnchor( - owner=value, - storage=storage, - allocation_address=storage.data_ptr(), - allocation_nbytes=storage.nbytes(), - ) - return _checked_memory( - registration_id, - name, - anchor, - address, - nbytes, - device, - memory_type, - ) - - untyped_storage_type = getattr(torch, "UntypedStorage", ()) - typed_storage_type = getattr(getattr(torch, "storage", object), "TypedStorage", ()) - if isinstance(value, (untyped_storage_type, typed_storage_type)): - storage = getattr(value, "_untyped_storage", value) - nbytes = storage.nbytes() - address = storage.data_ptr() - device = str(storage.device) - if storage.device.type == "cuda": - memory_type = "cuda" - else: - try: - memory_type = "host-pinned" if _storage_is_pinned(storage) else "host" - except RuntimeError: - memory_type = "host" - anchor = _LifetimeAnchor( - owner=value, - storage=storage, - allocation_address=address, - allocation_nbytes=nbytes, - ) - return _checked_memory( - registration_id, - name, - anchor, - address, - nbytes, - device, - memory_type, - ) - - try: - view = memoryview(value) # type: ignore[arg-type] - except TypeError as error: - raise InvalidRegionError( - "registration expects a contiguous Tensor, Storage, writable " - "buffer, or RawSpan" - ) from error - if not view.c_contiguous: - raise InvalidRegionError( - "buffer registration requires C-contiguous memory; " - "no implicit copy is performed" - ) - if view.readonly: - raise InvalidRegionError("buffer registration requires writable memory") - try: - address = ctypes.addressof(ctypes.c_char.from_buffer(view)) - except (TypeError, ValueError) as error: - raise InvalidRegionError("could not obtain the buffer address") from error - anchor = _LifetimeAnchor( - owner=value, - buffer=view, - allocation_address=address, - allocation_nbytes=view.nbytes, - ) - return _checked_memory( - registration_id, - name, - anchor, - address, - view.nbytes, - "cpu", - "host", - ) - - -def _checked_memory( - registration_id: str, - name: str, - owner: object, - address: int, - nbytes: int, - device: str, - memory_type: str, -) -> RegisteredMemory: - if not isinstance(address, int) or address <= 0 or nbytes <= 0: - raise InvalidRegionError("cannot register empty or unallocated memory") - return RegisteredMemory( - registration_id=registration_id, - name=name, - owner=owner, - address=address, - nbytes=nbytes, - device=device, - memory_type=memory_type, - ) - - -def _released_memory(memory: RegisteredMemory) -> RegisteredMemory: - return RegisteredMemory( - registration_id=memory.registration_id, - name=memory.name, - owner=None, - address=memory.address, - nbytes=memory.nbytes, - device=memory.device, - memory_type=memory.memory_type, - ) - - -def _memory_ranges_overlap(first: RegisteredMemory, second: RegisteredMemory) -> bool: - if first.device != second.device: - return False - return ( - first.address < second.address + second.nbytes - and second.address < first.address + first.nbytes - ) - - -def _validate_memory_stable(memory: RegisteredMemory) -> None: - anchor = memory.owner - if not isinstance(anchor, _LifetimeAnchor): - return - if anchor.storage is not None: - storage = cast(Any, anchor.storage) - try: - address = storage.data_ptr() - nbytes = storage.nbytes() - except Exception as error: - raise InvalidRegionError("could not validate registered Storage") from error - if address != anchor.allocation_address or nbytes != anchor.allocation_nbytes: - raise InvalidRegionError( - "registered Storage was resized; deregister and register it again" - ) - if anchor.buffer is not None: - try: - address = ctypes.addressof(ctypes.c_char.from_buffer(anchor.buffer)) - except (TypeError, ValueError) as error: - raise InvalidRegionError( - "registered buffer no longer exposes its original address" - ) from error - if ( - address != anchor.allocation_address - or anchor.buffer.nbytes != anchor.allocation_nbytes - ): - raise InvalidRegionError("registered buffer address or size changed") - - -def _tensor_memory_type(tensor: torch.Tensor) -> str: - if tensor.device.type == "cuda": - return "cuda" - if tensor.device.type == "cpu": - try: - return "host-pinned" if tensor.is_pinned() else "host" - except RuntimeError: - return "host" - return tensor.device.type - - -def _storage_is_pinned(storage: object) -> bool: - view = torch.empty( - 0, - dtype=torch.uint8, - device=storage.device, # type: ignore[attr-defined] - ) - return view.set_(storage).is_pinned() # type: ignore[arg-type] - - -def _validate_region_bounds( - allocation_nbytes: int, - offset: int, - nbytes: int | None, - stride: int | None, - count: int | None, -) -> tuple[int, int, int | None, int | None]: - _validate_nonnegative_int(offset, "offset") - stride, count = _validate_stride(stride, count) - if nbytes is None: - if stride is not None: - raise InvalidRegionError( - "nbytes is required when stride and count are specified" - ) - nbytes = allocation_nbytes - offset - _validate_positive_int(nbytes, "nbytes") - if stride is not None and stride < nbytes: - raise InvalidRegionError("stride must be at least nbytes") - extent = _region_extent(nbytes, stride, count) - if offset + extent > allocation_nbytes: - raise InvalidRegionError( - f"region [{offset}, {offset + extent}) exceeds allocation size " - f"{allocation_nbytes}" - ) - return offset, nbytes, stride, count - - -def _validate_stride( - stride: int | None, - count: int | None, -) -> tuple[int | None, int | None]: - if (stride is None) != (count is None): - raise InvalidRegionError("stride and count must be specified together") - if stride is not None: - _validate_positive_int(stride, "stride") - _validate_positive_int(count, "count") - return stride, count - - -def _region_extent( - nbytes: int, - stride: int | None, - count: int | None, -) -> int: - if stride is None: - return nbytes - if count is None: - raise InvalidRegionError("count is required with stride") - return (count - 1) * stride + nbytes - - -def _validate_plan_indices( - plan: TransferPlan, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - op: TransferOp, -) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]: - if not plan.indexed: - if local_indices is not None or remote_indices is not None: - raise InvalidRegionError( - "indices require a plan prepared with indexed=True" - ) - local = remote = None - else: - if (local_indices is None) != (remote_indices is None): - raise InvalidRegionError( - "local_indices and remote_indices must be specified together" - ) - if local_indices is None: - local = remote = None - else: - remote_values = cast(Sequence[int], remote_indices) - local = _validate_indices( - local_indices, plan.local_block_count, "local_indices" - ) - remote = _validate_indices( - remote_values, - plan.remote_block_count, - "remote_indices", - ) - if local is None: - if plan._local_total_nbytes != plan._remote_total_nbytes: - raise InvalidRegionError( - f"selected local and remote byte totals differ: " - f"{plan._local_total_nbytes} != {plan._remote_total_nbytes}" - ) - destination_nonoverlap = ( - plan._remote_full_nonoverlap - if op is TransferOp.WRITE - else plan._local_full_nonoverlap - ) - if not destination_nonoverlap: - raise InvalidRegionError("selected destination regions must not overlap") - return local, remote - - selected_local = ( - None - if local is None - else _select_region_blocks(plan._local, local, plan._local_starts) - ) - selected_remote = ( - None - if remote is None - else _select_region_blocks(plan._remote, remote, plan._remote_starts) - ) - local_size = _selected_region_nbytes(plan._local, selected_local) - remote_size = _selected_region_nbytes(plan._remote, selected_remote) - if local_size != remote_size: - raise InvalidRegionError( - f"selected local and remote byte totals differ: " - f"{local_size} != {remote_size}" - ) - if op is TransferOp.WRITE: - _validate_nonoverlapping_destination(plan._remote, selected_remote) - else: - _validate_nonoverlapping_destination(plan._local, selected_local) - return local, remote - - -def _snapshot_backend_index_selection( - plan: TransferPlan, - op: TransferOp, - local_indices: Sequence[int], - remote_indices: Sequence[int], - *, - collect_local_devices: bool = True, -) -> tuple[BackendIndexSelection, tuple[str, ...]]: - """Create the sole private immutable selection from one input snapshot.""" - - ( - packed, - local_count, - remote_count, - selected_nbytes, - uniform_pair_layout, - destination_catalog_nonoverlapping, - local_devices, - ) = _pack_and_validate_index_selection( - plan, - op, - local_indices, - remote_indices, - collect_local_devices=collect_local_devices, - ) - # ``packed`` is owned immutable bytes, so the cast is already read-only. - # Avoid allocating a second memoryview wrapper on every dynamic selection. - packed_view = memoryview(packed).cast("i") - local_view = packed_view[:local_count] - remote_view = packed_view[local_count : local_count + remote_count] - return ( - _make_backend_index_selection( - op, - plan._handle, - local_view, - remote_view, - selected_nbytes=selected_nbytes, - uniform_pair_layout=uniform_pair_layout, - destination_catalog_nonoverlapping=destination_catalog_nonoverlapping, - destination_indices_unique=True, - ), - local_devices, - ) - - -def _pack_and_validate_index_selection( - plan: TransferPlan, - op: TransferOp, - local_indices: Sequence[int], - remote_indices: Sequence[int], - *, - collect_local_devices: bool = True, -) -> tuple[bytes, int, int, int, bool, bool, tuple[str, ...]]: - """Snapshot two vectors once and derive reusable Core proof facts.""" - - local_input_view = _native_int32_buffer_view(local_indices) - remote_input_view = _native_int32_buffer_view(remote_indices) - native_destination_duplicate: bool | None = None - if local_input_view is not None and remote_input_view is not None: - local_count = local_input_view.nbytes // 4 - if local_count == 0: - raise InvalidRegionError("local_indices must be non-empty") - remote_count = remote_input_view.nbytes // 4 - if remote_count == 0: - raise InvalidRegionError("remote_indices must be non-empty") - - if _compiled_index_analyzer is not None: - # The helper owns the only copy: it creates one combined PyBytes - # snapshot under the GIL, then derives every fact from that - # immutable snapshot with the GIL released. - ( - packed, - analyzed_local_count, - analyzed_remote_count, - local_minimum, - local_maximum, - remote_minimum, - remote_maximum, - native_destination_duplicate, - ) = _compiled_index_analyzer( - local_input_view, - remote_input_view, - op is TransferOp.WRITE, - ) - if ( - analyzed_local_count != local_count - or analyzed_remote_count != remote_count - ): - raise RuntimeError( - "compiled int32 index analyzer returned inconsistent counts" - ) - local_bounds = (local_minimum, local_maximum) - remote_bounds = (remote_minimum, remote_maximum) - else: - # Compatibility path for source-tree Python loaded with an older - # torch._C binary. It preserves the exact pre-helper behavior and - # is deliberately selected only by a missing import, never by a - # helper exception. - packed = b"".join((local_input_view.cast("B"), remote_input_view.cast("B"))) - packed_indices = memoryview(packed).cast("i") - local_view = packed_indices[:local_count] - remote_view = packed_indices[local_count : local_count + remote_count] - destination_view = remote_view if op is TransferOp.WRITE else local_view - if max(local_count, remote_count) <= _INDEX_SELECTION_TORCH_THRESHOLD: - local_bounds = (min(local_view), max(local_view)) - remote_bounds = (min(remote_view), max(remote_view)) - native_destination_duplicate = len(set(destination_view)) != len( - destination_view - ) - else: - validation_scratch = bytearray(packed) - validation = torch.frombuffer(validation_scratch, dtype=torch.int32) - local_tensor = validation[:local_count] - remote_tensor = validation[local_count : local_count + remote_count] - local_minimum, local_maximum = torch.aminmax(local_tensor) - remote_minimum, remote_maximum = torch.aminmax(remote_tensor) - destination_tensor = ( - remote_tensor if op is TransferOp.WRITE else local_tensor - ) - ordered_destination = torch.sort(destination_tensor).values - local_bounds = (int(local_minimum), int(local_maximum)) - remote_bounds = (int(remote_minimum), int(remote_maximum)) - native_destination_duplicate = bool( - torch.any(ordered_destination[1:] == ordered_destination[:-1]) - ) - - packed_indices = memoryview(packed).cast("i") - local_view = packed_indices[:local_count] - remote_view = packed_indices[local_count : local_count + remote_count] - - # Preserve the public error order: local bounds, remote bounds, byte - # equality, then destination duplicate/physical-overlap diagnostics. - _raise_packed_index_bounds( - *local_bounds, plan.local_block_count, "local_indices" - ) - _raise_packed_index_bounds( - *remote_bounds, plan.remote_block_count, "remote_indices" - ) - else: - # Mixed buffers and general Sequences retain the established - # operator.index normalization and validation behavior. - packed = bytearray() - local_count, local_needs_bounds = _append_index_snapshot( - packed, - local_indices, - plan.local_block_count, - "local_indices", - ) - if local_needs_bounds: - # Preserve the public left-to-right validation contract for mixed - # inputs. A native local buffer is copied by _append_index_snapshot, - # but its bounds must be diagnosed before the remote Sequence is - # iterated (which may itself be invalid or have side effects). - local_snapshot_view = memoryview(packed).cast("i") - try: - _validate_packed_index_bounds( - local_snapshot_view, - plan.local_block_count, - "local_indices", - ) - finally: - # bytearray cannot be resized while an exported view exists. - local_snapshot_view.release() - remote_count, remote_needs_bounds = _append_index_snapshot( - packed, - remote_indices, - plan.remote_block_count, - "remote_indices", - ) - packed_indices = memoryview(packed).cast("i") - local_view = packed_indices[:local_count] - remote_view = packed_indices[local_count : local_count + remote_count] - if remote_needs_bounds: - _validate_packed_index_bounds( - remote_view, plan.remote_block_count, "remote_indices" - ) - - selected_local: tuple[tuple[Region, int], ...] | None = None - selected_remote: tuple[tuple[Region, int], ...] | None = None - if plan._local_uniform_nbytes is None: - selected_local = _select_region_blocks( - plan._local, local_view, plan._local_starts - ) - local_size = _selected_region_nbytes(plan._local, selected_local) - else: - local_size = local_count * plan._local_uniform_nbytes - if plan._remote_uniform_nbytes is None: - selected_remote = _select_region_blocks( - plan._remote, remote_view, plan._remote_starts - ) - remote_size = _selected_region_nbytes(plan._remote, selected_remote) - else: - remote_size = remote_count * plan._remote_uniform_nbytes - if local_size != remote_size: - raise InvalidRegionError( - f"selected local and remote byte totals differ: " - f"{local_size} != {remote_size}" - ) - - destination_full_nonoverlap = ( - plan._remote_full_nonoverlap - if op is TransferOp.WRITE - else plan._local_full_nonoverlap - ) - if native_destination_duplicate: - raise InvalidRegionError("selected destination regions must not overlap") - if destination_full_nonoverlap: - destination_view = remote_view if op is TransferOp.WRITE else local_view - if native_destination_duplicate is None: - _validate_unique_packed_indices(destination_view) - elif op is TransferOp.WRITE: - if selected_remote is None: - selected_remote = _select_region_blocks( - plan._remote, remote_view, plan._remote_starts - ) - _validate_nonoverlapping_destination(plan._remote, selected_remote) - else: - if selected_local is None: - selected_local = _select_region_blocks( - plan._local, local_view, plan._local_starts - ) - _validate_nonoverlapping_destination(plan._local, selected_local) - - if not collect_local_devices: - local_devices = () - elif plan._local_uniform_device is None: - if selected_local is None: - selected_local = _select_region_blocks( - plan._local, local_view, plan._local_starts - ) - local_devices = tuple( - dict.fromkeys(region.device for region, _ in selected_local) - ) - else: - local_devices = (plan._local_uniform_device,) - - uniform_pair_layout = ( - local_count == remote_count - and plan._local_uniform_nbytes is not None - and plan._local_uniform_nbytes == plan._remote_uniform_nbytes - ) - return ( - bytes(packed), - local_count, - remote_count, - local_size, - uniform_pair_layout, - destination_full_nonoverlap, - local_devices, - ) - - -def _append_index_snapshot( - packed: bytearray, - values: Sequence[int], - upper_bound: int, - name: str, -) -> tuple[int, bool]: - """Append one vector, combining general-sequence coercion and bounds.""" - - if values is None: - raise InvalidRegionError(f"{name} must be non-empty") - view = _native_int32_buffer_view(values) - if view is not None: - count = view.nbytes // 4 - if count == 0: - raise InvalidRegionError(f"{name} must be non-empty") - packed.extend(view.cast("B")) - return count, True - - count = 0 - try: - iterator = iter(values) - except TypeError as error: - raise InvalidRegionError(f"{name} must contain integers") from error - for raw_value in iterator: - if isinstance(raw_value, bool): - raise InvalidRegionError(f"{name} must contain integers") - try: - value = operator.index(raw_value) - except TypeError as error: - raise InvalidRegionError(f"{name} must contain integers") from error - if value < 0 or value >= upper_bound: - raise InvalidRegionError( - f"{name} index {value} is outside [0, {upper_bound})" - ) - if value > _INT32_MAX: - raise InvalidRegionError( - f"{name} index {value} exceeds the packed int32 range" - ) - packed.extend(value.to_bytes(4, byteorder=sys.byteorder, signed=True)) - count += 1 - if count == 0: - raise InvalidRegionError(f"{name} must be non-empty") - return count, False - - -def _native_int32_buffer_view(values: object) -> memoryview | None: - """Return the exact packed-input view, without copying or coercion.""" - - try: - view = memoryview(values) - except TypeError: - return None - if ( - view.ndim == 1 - and view.c_contiguous - and view.itemsize == 4 - and view.format in _NATIVE_INT32_FORMATS - ): - return view - return None - - -def _raise_packed_index_bounds( - low: int, - high: int, - upper_bound: int, - name: str, -) -> None: - if low < 0: - raise InvalidRegionError(f"{name} index {low} is outside [0, {upper_bound})") - if high >= upper_bound: - raise InvalidRegionError(f"{name} index {high} is outside [0, {upper_bound})") - - -def _validate_packed_index_bounds( - values: memoryview, - upper_bound: int, - name: str, -) -> None: - tensor = torch.frombuffer(values, dtype=torch.int32) - minimum, maximum = torch.aminmax(tensor) - low = int(minimum) - high = int(maximum) - _raise_packed_index_bounds(low, high, upper_bound, name) - - -def _validate_unique_packed_indices(values: memoryview) -> None: - if len(values) < 2: - return - tensor = torch.frombuffer(values, dtype=torch.int32) - ordered = torch.sort(tensor).values - if bool(torch.any(ordered[1:] == ordered[:-1])): - raise InvalidRegionError("selected destination regions must not overlap") - - -def _validate_cuda_ordering_arguments( - capabilities: Capabilities, - backend: str, - mode: CudaOrderingMode, - stream: object | None, - ready_event: object | None, -) -> None: - if stream is not None and ready_event is not None: - raise ValueError("stream and ready_event are mutually exclusive") - ordering_token_count = int(stream is not None) + int(ready_event is not None) - if mode is CudaOrderingMode.CALLER_READY: - if ordering_token_count: - raise ValueError( - "caller_ready CUDA ordering does not accept a stream or ready_event" - ) - elif ordering_token_count != 1: - raise ValueError( - "host_wait CUDA ordering requires exactly one stream or ready_event" - ) - if mode not in capabilities.cuda_ordering_modes: - raise UnsupportedError( - f"backend {backend!r} does not support CUDA ordering " - f"mode {mode.value!r}" - ) - - -def _validate_attached_notification( - capabilities: Capabilities, - backend: str, - notification: bytes | bytearray | memoryview | None, -) -> bytes | None: - if notification is None: - return None - if not capabilities.attached_notifications: - raise UnsupportedError(f"backend {backend!r} does not support notifications") - result = _as_bytes( - notification, - "notification", - max_nbytes=_MAX_NOTIFICATION_BYTES, - ) - if not result: - raise ValueError("an attached notification payload must be non-empty") - return result - - -def _validate_request_slot_notification( - notification: bytes | bytearray | memoryview | None, -) -> bytes | None: - """Validate an explicit per-generation slot notification or clear.""" - - if notification is None: - return None - result = _as_bytes( - notification, - "notification", - max_nbytes=_MAX_NOTIFICATION_BYTES, - ) - if not result: - raise ValueError("an attached notification payload must be non-empty") - return result - - -def _validate_cuda_ordering_selection( - plan: TransferPlan, - local_indices: Sequence[int] | None, - mode: CudaOrderingMode, - stream: object | None, - ready_event: object | None, -) -> torch.device | None: - """Validate a CUDA submission shape without establishing readiness.""" - - if mode is CudaOrderingMode.CALLER_READY: - return None - if mode is not CudaOrderingMode.HOST_WAIT: - raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") - - selected = ( - tuple((region, 0) for region in plan._local) - if local_indices is None - else _select_region_blocks(plan._local, local_indices, plan._local_starts) - ) - return _validate_cuda_ordering_devices( - tuple(dict.fromkeys(region.device for region, _ in selected)), - mode, - stream, - ready_event, - ) - - -def _validate_cuda_ordering_devices( - selected_devices: Sequence[str], - mode: CudaOrderingMode, - stream: object | None, - ready_event: object | None, -) -> torch.device | None: - """Validate ordering from device facts cached by an IndexSelection.""" - - if mode is CudaOrderingMode.CALLER_READY: - return None - if mode is not CudaOrderingMode.HOST_WAIT: - raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") - - devices: set[torch.device] = set() - for device_name in selected_devices: - try: - device = torch.device(device_name) - except (RuntimeError, TypeError, ValueError) as error: - raise InvalidRegionError( - f"local region has invalid device {device_name!r}" - ) from error - if device.type != "cuda" or device.index is None: - raise InvalidRegionError( - "host_wait requires selected local regions on one explicitly " - "indexed CUDA device" - ) - devices.add(device) - if len(devices) != 1: - raise UnsupportedError( - "host_wait requires selected local regions on exactly one CUDA device" - ) - expected_device = next(iter(devices)) - - if stream is not None: - stream_types = tuple( - stream_type - for stream_type in ( - getattr(torch, "Stream", None), - getattr(torch.cuda, "Stream", None), - ) - if isinstance(stream_type, type) - ) - if not stream_types or not isinstance(stream, stream_types): - raise TypeError("stream must be a torch CUDA Stream") - _validate_cuda_token_device(stream, expected_device, "stream") - else: - event_type = getattr(torch.cuda, "Event", None) - if not isinstance(event_type, type) or not isinstance(ready_event, event_type): - raise TypeError("ready_event must be a torch.cuda.Event") - _validate_cuda_token_device(ready_event, expected_device, "ready_event") - - return expected_device - - -def _establish_cuda_ordering_prevalidated( - mode: CudaOrderingMode, - stream: object | None, - ready_event: object | None, - expected_device: torch.device | None, -) -> None: - """Resolve Core-owned CUDA readiness before entering provider locks.""" - - if mode is CudaOrderingMode.CALLER_READY: - return - if mode is not CudaOrderingMode.HOST_WAIT or expected_device is None: - raise UnsupportedError(f"unsupported CUDA ordering mode {mode!r}") - - if stream is not None: - try: - if _cuda_token_stream_is_capturing(stream): - raise UnsupportedError("host_wait is not CUDA Graph capture safe") - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to query CUDA Graph capture state" - ) from error - - if stream is not None: - try: - event = stream.record_event() # type: ignore[attr-defined] - except Exception as error: - raise BackendFailureError( - "failed to record the CUDA readiness event" - ) from error - else: - event = ready_event - - if stream is not None: - # A stream implementation supplies the event on each submission; keep - # this runtime check even though the stream itself was validated once. - event_type = getattr(torch.cuda, "Event", None) - if not isinstance(event_type, type) or not isinstance(event, event_type): - raise TypeError("ready_event must be a torch.cuda.Event") - _validate_cuda_token_device(event, expected_device, "ready_event") - if event is None: - raise BackendFailureError("CUDA readiness event is missing") - try: - event.synchronize() # type: ignore[union-attr] - except TransferError: - raise - except Exception as error: - raise BackendFailureError("failed to wait for CUDA readiness") from error - - -def _cuda_token_stream_is_capturing(stream: object) -> bool: - """Query the supplied stream before Core records a new readiness event.""" - - direct_query = getattr(stream, "is_capturing", None) - if callable(direct_query): - return bool(direct_query()) - with torch.cuda.stream(stream): # type: ignore[arg-type] - return torch.cuda.is_current_stream_capturing() - - -def _make_reusable_cuda_stream_readiness( - stream: object, - expected_device: torch.device | None, -) -> Callable[[], None]: - """Create one per-slot event and return its generation readiness gate.""" - - if expected_device is None: - raise UnsupportedError("host_wait stream requires one CUDA device") - try: - event = torch.cuda.Event() - record = cast(Callable[[object], None], event.record) - synchronize = cast(Callable[[], None], event.synchronize) - except Exception as error: - raise BackendFailureError( - "failed to create the reusable CUDA readiness event" - ) from error - - direct_capture_query = getattr(stream, "is_capturing", None) - if callable(direct_capture_query): - capture_query = direct_capture_query - else: - stream_context = torch.cuda.stream - current_stream_is_capturing = torch.cuda.is_current_stream_capturing - - def capture_query() -> bool: - with stream_context(stream): # type: ignore[arg-type] - return bool(current_stream_is_capturing()) - - device_validated = False - - def synchronize_generation() -> None: - nonlocal device_validated - try: - if capture_query(): - raise UnsupportedError("host_wait is not CUDA Graph capture safe") - except TransferError: - raise - except Exception as error: - raise BackendFailureError( - "failed to query CUDA Graph capture state" - ) from error - try: - record(stream) - except Exception as error: - raise BackendFailureError( - "failed to record the CUDA readiness event" - ) from error - if not device_validated: - _validate_cuda_token_device(event, expected_device, "ready_event") - device_validated = True - try: - synchronize() - except TransferError: - raise - except Exception as error: - raise BackendFailureError("failed to wait for CUDA readiness") from error - - return synchronize_generation - - -def _validate_cuda_token_device( - token: object, - expected: torch.device, - name: str, -) -> None: - device_value = getattr(token, "device", None) - if device_value is None: - qualifier = "recorded " if name == "ready_event" else "" - raise InvalidRegionError(f"{name} must be a {qualifier}CUDA token") - try: - device = torch.device(device_value) - except (RuntimeError, TypeError, ValueError) as error: - raise InvalidRegionError(f"{name} has an invalid device") from error - if device != expected: - raise InvalidRegionError( - f"{name} device {device} does not match local region device {expected}" - ) - - -def _logical_block_count(regions: Sequence[Region]) -> int: - return sum(region.count or 1 for region in regions) - - -def _uniform_region_nbytes(regions: Sequence[Region]) -> int | None: - first = regions[0].nbytes - return first if all(region.nbytes == first for region in regions[1:]) else None - - -def _uniform_region_device(regions: Sequence[Region]) -> str | None: - first = regions[0].device - return first if all(region.device == first for region in regions[1:]) else None - - -def _catalog_starts(regions: Sequence[Region]) -> tuple[int, ...]: - starts: list[int] = [] - next_index = 0 - for region in regions: - starts.append(next_index) - next_index += region.count or 1 - return tuple(starts) - - -def _select_region_blocks( - regions: Sequence[Region], - indices: Sequence[int], - starts: Sequence[int], -) -> tuple[tuple[Region, int], ...]: - selected: list[tuple[Region, int]] = [] - for index in indices: - region_index = bisect_right(starts, index) - 1 - if region_index < 0 or region_index >= len(regions): - raise InvalidRegionError( - "logical block index is outside the region catalog" - ) - selected.append((regions[region_index], index - starts[region_index])) - return tuple(selected) - - -def _selected_region_nbytes( - regions: Sequence[Region], - selected: Sequence[tuple[Region, int]] | None, -) -> int: - if selected is None: - return sum(region.total_nbytes for region in regions) - return sum(region.nbytes for region, _ in selected) - - -def _validate_nonoverlapping_destination( - regions: Sequence[Region], - selected: Sequence[tuple[Region, int]] | None, -) -> None: - intervals: dict[object, list[tuple[int, int]]] = {} - blocks = ( - ( - (region, block_index) - for region in regions - for block_index in range(region.count or 1) - ) - if selected is None - else iter(selected) - ) - for region, block_index in blocks: - descriptor = region._backend_descriptor - if isinstance(region._origin, Registration): - key: object = ("local", region.device) - base = region._origin._memory.address - elif isinstance(descriptor, RegistrationMetadata): - key = ( - "remote-registration", - id(region._origin), - descriptor.registration_id, - ) - base = 0 - elif isinstance(descriptor, _RemoteRawDescriptor): - key = ( - "remote-address", - id(region._origin), - descriptor.device, - descriptor.memory_type, - ) - base = descriptor.address - else: - raise InvalidRegionError("region has an invalid memory descriptor") - stride = region.stride or region.nbytes - start = base + region.offset + block_index * stride - intervals.setdefault(key, []).append((start, start + region.nbytes)) - for values in intervals.values(): - values.sort() - for previous, current in itertools.pairwise(values): - if current[0] < previous[1]: - raise InvalidRegionError( - "selected destination regions must not overlap" - ) - - -def _regions_are_nonoverlapping(regions: Sequence[Region]) -> bool: - # Each compact Region is internally disjoint because stride >= nbytes. - # Compare only its bounding envelope against other regions in the same - # physical address space. An ambiguous interleaving is conservatively - # rejected instead of expanding `count` logical blocks. - envelopes: dict[object, list[tuple[int, int]]] = {} - for region in regions: - descriptor = region._backend_descriptor - if isinstance(region._origin, Registration): - key: object = ("local", region.device) - base = region._origin._memory.address - elif isinstance(descriptor, RegistrationMetadata): - key = ( - "remote-registration", - id(region._origin), - descriptor.registration_id, - ) - base = 0 - elif isinstance(descriptor, _RemoteRawDescriptor): - key = ( - "remote-address", - id(region._origin), - descriptor.device, - descriptor.memory_type, - ) - base = descriptor.address - else: - return False - stride = region.stride or region.nbytes - start = base + region.offset - end = start + ((region.count or 1) - 1) * stride + region.nbytes - envelopes.setdefault(key, []).append((start, end)) - for values in envelopes.values(): - values.sort() - for previous, current in itertools.pairwise(values): - if current[0] < previous[1]: - return False - return True - - -def _validate_indices( - values: Sequence[int], - upper_bound: int, - name: str, -) -> tuple[int, ...]: - result: list[int] = [] - for raw_value in values: - if isinstance(raw_value, bool): - raise InvalidRegionError(f"{name} must contain integers") - try: - value = operator.index(raw_value) - except TypeError as error: - raise InvalidRegionError(f"{name} must contain integers") from error - if value < 0 or value >= upper_bound: - raise InvalidRegionError( - f"{name} index {value} is outside [0, {upper_bound})" - ) - result.append(value) - if not result: - raise InvalidRegionError(f"{name} must be non-empty") - return tuple(result) - - -def _adopt_resource_handle(resource: object, handle: object, kind: str) -> None: - if handle is None or handle is _UNADOPTED_HANDLE: - raise BackendFailureError(f"backend adopted an invalid {kind} handle") - mutable_resource = cast(Any, resource) - current = mutable_resource._handle - if current is _UNADOPTED_HANDLE: - mutable_resource._handle = handle - return - if current is handle: - return - raise BackendFailureError(f"backend attempted to adopt {kind} more than once") - - -def _creation_cleanup_failure( - kind: str, - original_error: BaseException, - cleanup_error: BaseException, -) -> BackendFailureError: - del cleanup_error - failure = BackendFailureError( - f"{kind} creation failed and cleanup is pending; recover it from the " - f"Endpoint open-resource snapshots and retry close()" - ) - add_note = getattr(failure, "add_note", None) - if callable(add_note): - add_note(f"original failure type: {type(original_error).__name__}") - return failure - - -def _endpoint_creation_cleanup_failure( - endpoint: Endpoint, - original_error: BaseException, - cleanup_error: BaseException, -) -> BackendFailureError: - failure = BackendFailureError( - "endpoint backend initialization failed and cleanup is pending; recover " - "exception.recovery_endpoint and retry close()" - ) - failure.recovery_endpoint = endpoint # type: ignore[attr-defined] - add_note = getattr(failure, "add_note", None) - if callable(add_note): - add_note(f"initialization failure type: {type(original_error).__name__}") - add_note(f"cleanup failure type: {type(cleanup_error).__name__}") - return failure - - -def _as_bytes( - value: object, - name: str, - *, - max_nbytes: int | None = None, -) -> bytes: - if type(value) is bytes: - if max_nbytes is not None and len(value) > max_nbytes: - raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") - return value - if not isinstance(value, (bytes, bytearray, memoryview)): - raise TypeError(f"{name} must be bytes-like") - if max_nbytes is not None: - try: - source_nbytes = ( - value.nbytes if isinstance(value, memoryview) else len(value) - ) - except ValueError as error: - raise ValueError(f"{name} buffer is unavailable") from error - if source_nbytes > max_nbytes: - raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") - result = bytes(value) - if max_nbytes is not None and len(result) > max_nbytes: - raise ValueError(f"{name} exceeds the {max_nbytes}-byte size limit") - return result - - -def _validate_positive_int(value: object, name: str) -> None: - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise InvalidRegionError(f"{name} must be a positive integer") - - -def _transfer_slot_deadline_ns(timeout: float | None) -> int | None: - if timeout is None: - return None - if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): - raise TypeError("timeout must be a non-negative number or None") - if timeout < 0 or (isinstance(timeout, float) and not math.isfinite(timeout)): - raise ValueError("timeout must be finite and non-negative") - if isinstance(timeout, int): - timeout_ns = timeout * 1_000_000_000 - else: - numerator, denominator = timeout.as_integer_ratio() - timeout_ns = (numerator * 1_000_000_000 + denominator - 1) // denominator - return time.monotonic_ns() + timeout_ns - - -def _validate_bounded_poll_arguments( - *, - max_polls: object, - timeout_ns: object, - timeout_check_interval: object | None, - allow_zero: bool, -) -> None: - minimum = 0 if allow_zero else 1 - if isinstance(max_polls, bool) or not isinstance(max_polls, int): - qualifier = "non-negative" if allow_zero else "positive" - raise TypeError(f"max_polls must be a {qualifier} integer") - if max_polls < minimum: - qualifier = "non-negative" if allow_zero else "positive" - raise ValueError(f"max_polls must be a {qualifier} integer") - if max_polls > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS: - qualifier = "non-negative" if allow_zero else "positive" - raise ValueError(f"max_polls must be a {qualifier} signed 64-bit integer") - if timeout_ns is not None and ( - isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int) - ): - raise TypeError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - if timeout_ns is not None and ( - timeout_ns < 0 or timeout_ns > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS - ): - raise ValueError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - if timeout_check_interval is not None and ( - isinstance(timeout_check_interval, bool) - or not isinstance(timeout_check_interval, int) - ): - raise TypeError("timeout_check_interval must be a positive integer") - if timeout_check_interval is not None and timeout_check_interval <= 0: - raise ValueError("timeout_check_interval must be a positive integer") - if ( - timeout_check_interval is not None - and timeout_check_interval > _MAX_FUSED_REQUEST_SLOT_TIMEOUT_NS - ): - raise ValueError( - "timeout_check_interval must be a positive signed 64-bit integer" - ) - - -def _validate_nonnegative_int(value: object, name: str) -> None: - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise InvalidRegionError(f"{name} must be a non-negative integer") - - -def _validate_nonempty_str( - value: object, - name: str, - error_type: type[Exception] = ValueError, -) -> None: - if not isinstance(value, str) or not value: - raise error_type(f"{name} must be a non-empty string") - - -def _coerce_enum(value: object, enum_type: type[Any], name: str) -> Any: - if isinstance(value, enum_type): - return value - try: - return enum_type(value) - except (TypeError, ValueError) as error: - choices = ", ".join(repr(item.value) for item in enum_type) - raise ValueError(f"{name} must be one of {choices}") from error + if self.poll() != "done": + raise RuntimeError("work is not DONE; keep polling or retain memory") + ep._backend.release_work(self._handle) + ep._works.remove(self) + self._closed = True + if not ep._works: + _live_endpoints.discard(ep) diff --git a/torch/distributed/_transfer/_backend.py b/torch/distributed/_transfer/_backend.py index 49a77b0e737e1..c7c900569d491 100644 --- a/torch/distributed/_transfer/_backend.py +++ b/torch/distributed/_transfer/_backend.py @@ -1,1034 +1,35 @@ -from __future__ import annotations - -import threading -import time -from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass, field from typing import Any, Protocol -from ._errors import BackendNotFoundError, BusyError, UnsupportedError -from ._types import ( - BackendNotification, - BackendNotificationBatch, - BackendRegion, - Capabilities, - EndpointMetadata, - ProgressMode, - ThreadMode, - TransferOp, - WorkState, -) - - -__all__ = [ - "BACKEND_FACTORY_API_VERSION", - "Backend", - "BackendFactory", - "BackendFactoryV2", - "BackendIndexSelection", - "BackendRequestSlot", - "BackendWork", - "ImportedRawMemory", - "ImportedRegistration", - "RegisteredMemory", - "get_backend_factory", - "register_backend", - "unregister_backend", -] - -_REQUEST_SLOT_TERMINAL_STATES = frozenset( - {WorkState.COMPLETED, WorkState.FAILED, WorkState.CANCELLED} -) -_SIGNED_INT64_MAX = (1 << 63) - 1 -_INDEX_SELECTION_TOKEN = object() - - -class BackendIndexSelection: - """Core-authenticated, immutable indexed-selection lowering input. - - Instances are constructed only by Core after it has snapshotted and - validated both index vectors against one exact plan and operation. The - read-only memoryviews share one private, Core-owned native-int32 buffer. - Providers must validate the exact type and plan/operation binding before - using the proof flags; :meth:`Backend.submit_prevalidated` does that for - source-compatible providers. - """ - - __slots__ = ( - "_destination_catalog_nonoverlapping", - "_destination_indices_unique", - "_local_indices", - "_op", - "_plan_handle", - "_remote_indices", - "_sealed", - "_selected_nbytes", - "_uniform_pair_layout", - ) - - def __init__( - self, - op: TransferOp, - plan_handle: object, - local_indices: memoryview, - remote_indices: memoryview, - *, - selected_nbytes: int, - uniform_pair_layout: bool, - destination_catalog_nonoverlapping: bool, - destination_indices_unique: bool, - _token: object, - ) -> None: - if _token is not _INDEX_SELECTION_TOKEN: - raise TypeError("BackendIndexSelection objects are created by Core") - # This object is created on every full-dynamic submission. Calling its - # rejecting Python __setattr__ for each field would repeat getattr and - # dispatch overhead nine times before the object is sealed. - set_attribute = object.__setattr__ - set_attribute(self, "_op", op) - set_attribute(self, "_plan_handle", plan_handle) - set_attribute(self, "_local_indices", local_indices) - set_attribute(self, "_remote_indices", remote_indices) - set_attribute(self, "_selected_nbytes", selected_nbytes) - set_attribute(self, "_uniform_pair_layout", uniform_pair_layout) - set_attribute( - self, - "_destination_catalog_nonoverlapping", - destination_catalog_nonoverlapping, - ) - set_attribute(self, "_destination_indices_unique", destination_indices_unique) - set_attribute(self, "_sealed", True) - - def __setattr__(self, name: str, value: object) -> None: - raise AttributeError("BackendIndexSelection is immutable") - - def __delattr__(self, name: str) -> None: - raise AttributeError("BackendIndexSelection is immutable") - - @property - def local_indices(self) -> memoryview: - return self._local_indices - - @property - def remote_indices(self) -> memoryview: - return self._remote_indices - - @property - def selected_nbytes(self) -> int: - return self._selected_nbytes - - @property - def uniform_pair_layout(self) -> bool: - """Whether Core proved equal, uniform per-pair block geometry.""" - - return self._uniform_pair_layout - - @property - def destination_catalog_nonoverlapping(self) -> bool: - """Whether Core proved that the full destination catalog is disjoint.""" - - return self._destination_catalog_nonoverlapping - - @property - def destination_indices_unique(self) -> bool: - """Whether Core proved the destination logical indices are unique.""" - - return self._destination_indices_unique - - def _validate_binding(self, op: TransferOp, plan_handle: object) -> None: - if type(self) is not BackendIndexSelection: - raise TypeError("selection must be a Core BackendIndexSelection") - if self._op is not op or self._plan_handle is not plan_handle: - raise ValueError("prevalidated selection does not match plan and op") - - -def _make_backend_index_selection( - op: TransferOp, - plan_handle: object, - local_indices: memoryview, - remote_indices: memoryview, - *, - selected_nbytes: int, - uniform_pair_layout: bool, - destination_catalog_nonoverlapping: bool, - destination_indices_unique: bool, -) -> BackendIndexSelection: - return BackendIndexSelection( - op, - plan_handle, - local_indices, - remote_indices, - selected_nbytes=selected_nbytes, - uniform_pair_layout=uniform_pair_layout, - destination_catalog_nonoverlapping=destination_catalog_nonoverlapping, - destination_indices_unique=destination_indices_unique, - _token=_INDEX_SELECTION_TOKEN, - ) - - -@dataclass(frozen=True, slots=True) -class RegisteredMemory: - """A local allocation whose lifetime is retained by PyTorch Core.""" - - registration_id: str - name: str - owner: object = field(repr=False) - address: int - nbytes: int - device: str - memory_type: str - - -@dataclass(frozen=True, slots=True) -class ImportedRegistration: - """A named remote allocation imported from endpoint metadata.""" - - peer_handle: object = field(repr=False) - registration_id: str - name: str - nbytes: int - device: str - memory_type: str - descriptor: bytes = field(repr=False) - - -@dataclass(frozen=True, slots=True) -class ImportedRawMemory: - """An expert-only remote span interpreted and validated by a backend.""" - - peer_handle: object = field(repr=False) - address: int - nbytes: int - device: str - memory_type: str - - -class BackendWork(ABC): - """Provider request used by the stable :class:`Work` wrapper. - - ``close`` must be idempotent. Core retains every submission argument until - the public Work is closed. Core also retains any CUDA readiness token even - though it resolves readiness before calling the provider. State may remain - PENDING, advance to RUNNING, or advance directly to a terminal state; - RUNNING may only remain RUNNING or advance to a terminal state. - """ - - # An empty base layout lets latency-sensitive provider subclasses opt into - # real slots. A subclass that omits ``__slots__`` still receives its normal - # instance dictionary, preserving the source-compatible extension model. - __slots__ = () - - @property - @abstractmethod - def state(self) -> WorkState | None: ... - - @property - @abstractmethod - def error(self) -> BaseException | None: ... - - @property - def completion_event(self) -> object | None: - return None - - def poll_bounded( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState: - """Perform at most ``max_polls`` state observations. - - Providers advertising ``capabilities.fused_work_poll`` should override - this method so the bounded loop remains inside one provider boundary. - The default preserves source compatibility with existing providers. - A zero timeout still permits the first observation; the timeout is - checked only between observations and cannot preempt a blocking poll. - """ - - if isinstance(max_polls, bool) or not isinstance(max_polls, int): - raise TypeError("max_polls must be a positive integer") - if max_polls <= 0 or max_polls > _SIGNED_INT64_MAX: - raise ValueError("max_polls must be a positive signed 64-bit integer") - if timeout_ns is not None: - if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): - raise TypeError("timeout_ns must be a non-negative integer or None") - if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - if isinstance(timeout_check_interval, bool) or not isinstance( - timeout_check_interval, int - ): - raise TypeError("timeout_check_interval must be a positive integer") - if timeout_check_interval <= 0 or timeout_check_interval > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_check_interval must be a positive signed 64-bit integer" - ) - deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns - state = WorkState.PENDING - for poll_index in range(max_polls): - state = self.state - if ( - not isinstance(state, WorkState) - or state in _REQUEST_SLOT_TERMINAL_STATES - ): - break - if ( - deadline_ns is not None - and (poll_index == 0 or (poll_index + 1) % timeout_check_interval == 0) - and time.monotonic_ns() >= deadline_ns - ): - break - return state - - @abstractmethod - def cancel(self) -> bool: - """Request best-effort cancellation. - - Returns ``True`` only when the request is known to be cancelled. - """ - - @abstractmethod - def close(self) -> None: - """Release the provider request after it reaches a terminal state.""" - - -class BackendRequestSlot(ABC): - """Provider-owned persistent request behind a public ``TransferSlot``. - - Core queries ``state``, ``error``, and ``completion_event`` only while one - unrecycled generation is active. ``start`` rearms the immutable transfer - described at construction. It must return ``None`` and, once native work - may have launched, must convert every launch failure (including a caught - ``BaseException``) into observable terminal slot state rather than raise. - - ``recycle`` is called only after terminal state and returns the request to - provider-idle state without changing its immutable transfer description. - It must return ``None`` and be idempotent after successful recycle so Core - can retry if an asynchronous interruption lands after the provider state - transition. Ordinarily, ``close`` is called only while provider-idle and - must likewise be idempotent after successful release. - - A provider advertising - ``capabilities.deferred_successful_request_slot_recycle`` additionally - retains an observable, clean ``COMPLETED`` receipt when Core omits the - successful recycle in ``TransferSlot.execute``. Its next ``start``, - ``start_and_poll``, ``start_with_notification``, - ``start_and_poll_with_notification``, or ``close`` must consume that - receipt inside the same provider boundary. Explicit ``recycle`` remains - supported and idempotent. - Failure, cancellation, ambiguous launch, and nonterminal state must never - be consumed implicitly, and ``poll_state`` must continue to report - ``COMPLETED`` until the receipt is consumed. - - Core serializes calls to one slot under ``SERIALIZED`` and ``MULTIPLE``; - ``SINGLE`` supplies the same guarantee through its one-owner-thread - contract. Providers therefore need synchronization only for state shared - between different slots or other backend resources. Independent slots may - make progress concurrently under ``MULTIPLE``. Asynchronous interruption - is recoverable at every provider boundary, but request-slot methods are not - recursively reentrant from a Python signal handler while the interrupted - method is still unwinding. - """ - - # See BackendWork: provider slots can be dictionary-free without requiring - # existing unslotted third-party implementations to change. - __slots__ = () - - @property - @abstractmethod - def state(self) -> WorkState | None: ... - - def poll_state(self) -> WorkState | None: - """Return current state through a cacheable method dispatch. - - Providers with a latency-sensitive request path should override this - method directly. The default preserves source compatibility with - providers that implement only the ``state`` property. ``None`` means - the provider is still idle because an interrupted ``start`` never - entered or acknowledged a generation. Once a generation is - acknowledged, the provider must report a ``WorkState`` until recycle. - """ - - return self.state - - @property - @abstractmethod - def error(self) -> BaseException | None: ... - - @property - def completion_event(self) -> object | None: - return None - - @abstractmethod - def start(self) -> None: - """Start one generation of the immutable request.""" - - def start_with_notification(self, notification: bytes | None) -> None: - """Start with an explicit request-local notification override. - - ``None`` explicitly clears the notification for this generation; - non-``None`` payloads are non-empty immutable bytes already validated - by Core. Before returning, the provider must synchronously copy or - retain the payload for as long as native work can access it; Core need - not retain the argument after this boundary. Providers advertise this - optional operation through - ``request_slot_notification_overrides``. The default is deliberately - unsupported so an older provider can never silently retain a stale - notification from a prior repost. - """ - - del notification - raise UnsupportedError( - "backend request slot does not support notification overrides" - ) - - def start_and_poll( - self, - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - """Start one generation and perform a bounded initial state poll. - - This source-compatible default composes :meth:`start` and - :meth:`poll_state`. Providers advertising - ``capabilities.fused_request_slot_start_poll`` should override it so - native post and status polling remain inside one provider boundary. - - At most ``max_polls`` state observations may be made. Zero is a - post-only budget: fused providers should return the state implied by - the native post result without issuing a status query. The default - source-compatible composition returns the cached ``state`` after - ``start``; only an overriding fused provider can guarantee that reading - it performs no native probe. ``timeout_ns`` bounds elapsed time from - entry, checked between observations; it cannot preempt a blocking - provider operation. With a positive budget, one observation is made - after start even for a zero timeout. A successful implementation must - return a :class:`WorkState`. ``None`` retains its idle/unacknowledged - meaning so Core can recover an interrupted start, but is not a valid - successful acknowledgement of a generation. - """ - - if isinstance(max_polls, bool) or not isinstance(max_polls, int): - raise TypeError("max_polls must be a non-negative integer") - if max_polls < 0 or max_polls > _SIGNED_INT64_MAX: - raise ValueError("max_polls must be a non-negative signed 64-bit integer") - if timeout_ns is not None: - if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): - raise TypeError("timeout_ns must be a non-negative integer or None") - if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - start_result = self.start() - if start_result is not None: - raise TypeError("backend request-slot start must return None") - if max_polls == 0: - return self.state - deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns - state: WorkState | None = None - for _ in range(max_polls): - state = self.poll_state() - if ( - state is None - or not isinstance(state, WorkState) - or state in _REQUEST_SLOT_TERMINAL_STATES - or (deadline_ns is not None and time.monotonic_ns() >= deadline_ns) - ): - break - return state - - def start_and_poll_with_notification( - self, - notification: bytes | None, - *, - max_polls: int, - timeout_ns: int | None, - ) -> WorkState | None: - """Start with a notification override and perform bounded polling. - - This source-compatible default composes - :meth:`start_with_notification` and :meth:`poll_state`. Providers that - also advertise ``fused_request_slot_start_poll`` should override it so - notification selection, native post, and the initial status probes - remain in one provider boundary. - - The same copy-or-retain requirement as :meth:`start_with_notification` - applies before this method returns. A zero poll budget has the same - post-only semantics and default limitation as :meth:`start_and_poll`. - """ - - if isinstance(max_polls, bool) or not isinstance(max_polls, int): - raise TypeError("max_polls must be a non-negative integer") - if max_polls < 0 or max_polls > _SIGNED_INT64_MAX: - raise ValueError("max_polls must be a non-negative signed 64-bit integer") - if timeout_ns is not None: - if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): - raise TypeError("timeout_ns must be a non-negative integer or None") - if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - result = self.start_with_notification(notification) - if result is not None: - raise TypeError( - "backend request-slot start_with_notification must return None" - ) - if max_polls == 0: - return self.state - deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns - state: WorkState | None = None - for _ in range(max_polls): - state = self.poll_state() - if ( - state is None - or not isinstance(state, WorkState) - or state in _REQUEST_SLOT_TERMINAL_STATES - or (deadline_ns is not None and time.monotonic_ns() >= deadline_ns) - ): - break - return state - - def poll_bounded( - self, - *, - max_polls: int, - timeout_ns: int | None, - timeout_check_interval: int, - ) -> WorkState | None: - """Perform bounded state observations for an active generation. - - Providers advertising ``capabilities.fused_request_slot_poll`` should - override this source-compatible default so multiple native status - observations remain inside one provider boundary. The capability does - not require start/poll fusion. - - At most ``max_polls`` observations are made. ``timeout_ns`` measures a - soft monotonic elapsed-time bound from entry; it is checked after the - first observation and then every ``timeout_check_interval`` - observations. It cannot preempt a blocking provider call, and sampled - checking may overshoot by up to one sampling interval of provider - calls. A zero timeout still makes exactly one observation. ``None`` - disables clock checks. A returned ``None`` retains its provider-idle - meaning and is invalid after Core has observed an acknowledged active - generation. - """ - - if isinstance(max_polls, bool) or not isinstance(max_polls, int): - raise TypeError("max_polls must be a positive integer") - if max_polls <= 0 or max_polls > _SIGNED_INT64_MAX: - raise ValueError("max_polls must be a positive signed 64-bit integer") - if timeout_ns is not None: - if isinstance(timeout_ns, bool) or not isinstance(timeout_ns, int): - raise TypeError("timeout_ns must be a non-negative integer or None") - if timeout_ns < 0 or timeout_ns > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_ns must be a non-negative signed 64-bit integer or None" - ) - if isinstance(timeout_check_interval, bool) or not isinstance( - timeout_check_interval, int - ): - raise TypeError("timeout_check_interval must be a positive integer") - if timeout_check_interval <= 0 or timeout_check_interval > _SIGNED_INT64_MAX: - raise ValueError( - "timeout_check_interval must be a positive signed 64-bit integer" - ) - deadline_ns = None if timeout_ns is None else time.monotonic_ns() + timeout_ns - state: WorkState | None = None - for poll in range(max_polls): - state = self.poll_state() - if ( - state is None - or not isinstance(state, WorkState) - or state in _REQUEST_SLOT_TERMINAL_STATES - ): - break - observations = poll + 1 - if ( - deadline_ns is not None - and (observations == 1 or observations % timeout_check_interval == 0) - and time.monotonic_ns() >= deadline_ns - ): - break - return state - - @abstractmethod - def cancel(self) -> bool: - """Request best-effort cancellation of the active generation.""" - - @abstractmethod - def recycle(self) -> None: - """Return a terminal generation to retry-safe provider-idle state.""" - - @abstractmethod - def close(self) -> None: - """Release this provider request while it is idle.""" +class Backend(Protocol): + """Provider contract; methods are called under the Endpoint's host lock. -class Backend(ABC): - """Experimental provider SPI for one-sided endpoint transfer. - - Background progress is provider-owned. Release methods must be idempotent - after a successful release and must make partial cleanup retry-safe. - Handle-producing methods construct a provider transaction record, pass it - to ``adopt_handle`` before their first native mutation, and return ``None``. - Core can therefore roll back an adopted partial mutation or retain it for a - public close retry. Before adoption, a provider must fully roll back any - internal state before raising. If native cleanup fails, it may additionally - retain the orphan internally, reject conflicting operations, and retry - cleanup on later calls and close. + Regions are (address, byte_length, device_id) rows. prepare also accepts + (address, byte_length, device_id, stride, count) rows without expansion. + kind is 'host' or 'cuda'. Providers validate indices, ranges and sizes. + All handles are provider-private. poll/post return pending/done/error. + Core calls release_work only after done. Release failures retain ownership. + Metadata and notification bytes are opaque to Core, not authenticated. """ - @property - @abstractmethod - def capabilities(self) -> Capabilities: ... - - @property - def supported_thread_modes(self) -> frozenset[ThreadMode]: - return frozenset({ThreadMode.SERIALIZED}) - - @abstractmethod - def endpoint_payload(self, registration_handles: Sequence[object]) -> bytes: - """Export connection metadata limited to the selected registrations. - - A provider must not silently grant access to omitted registrations. It - may reject selective export if its transport cannot enforce it. - """ - - ... - - @abstractmethod - def register( - self, - memory: RegisteredMemory, - *, - adopt_handle: Callable[[object], None], - ) -> None: - """Adopt a retry-safe provider record before native registration.""" - - @abstractmethod - def deregister(self, registration_handle: object) -> None: ... - - @abstractmethod - def describe_registration(self, registration_handle: object) -> bytes: ... - - @abstractmethod - def import_peer( - self, - metadata: EndpointMetadata, - *, - adopt_handle: Callable[[object], None], - ) -> None: - """Adopt a retry-safe provider record before native peer mutation.""" - - @abstractmethod - def release_peer(self, peer_handle: object) -> None: ... - - @abstractmethod - def prepare( - self, - local: Sequence[BackendRegion], - remote: Sequence[BackendRegion], - *, - indexed: bool, - adopt_handle: Callable[[object], None], - ) -> None: - """Adopt a retry-safe provider record before native preparation.""" - - @abstractmethod - def release_plan(self, plan_handle: object) -> None: ... - - @abstractmethod - def submit( - self, - op: TransferOp, - plan_handle: object, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | None, - reuse_token: object | None, - adopt_work: Callable[[BackendWork], None], - ) -> None: - """Submit work after Core has established local CUDA readiness. - - Core resolves the selected CUDA ordering mode before entering the - serialized provider call, so this method must not add an implicit device - synchronization. Terminal completion means local memory may be safely - reused. ``completion_event``, when present, must express that completion - to the device. An attached notification may become visible only after the - remote writes are visible. - - Providers must also reject overlapping physical destination spans, - including aliases advertised as distinct remote registrations. - - For an indexed plan, each prepared ``BackendRegion`` contributes - ``count`` logical blocks, or one block when ``count`` is absent. Indices - address the flattened region-major, block-minor catalogs. A logical - block has ``nbytes`` bytes at ``offset + within_region_index * stride`` - (with an effective stride of ``nbytes`` for an unstrided region); - providers may keep this representation compact rather than expanding it - at prepare. ``None`` indices select every logical block in that catalog. - - ``reuse_token`` is an opaque identity hint. Core passes the same - non-``None`` object for every submission from one ``BoundTransfer`` and - for repeated full-selection, notification-free, caller-ready convenience - submissions of one operation on one plan. A provider may key immutable - cached lowering state by object identity, but must not inspect the token, - derive transfer semantics from it, or require it to be non-``None``. - Core tokens support weak references; caches should be weak-keyed or - scoped to and cleared with the provider plan. - - ``submit_prevalidated`` has a narrower token contract: its token - identifies the exact plan/operation/index selection, while notification - is request-local and may change. Only an override that can update the - native notification at post time may reuse that token. The default - therefore clears it when forwarding to this legacy method. - - Submission is transactional with respect to ownership. The provider - must construct its BackendWork and synchronously call ``adopt_work`` - before launching transport work. Before adoption it may raise only if - no transport access can remain in flight. After successful adoption it - must not raise: launch failures, including asynchronous interruptions, - must instead become observable through the adopted BackendWork. This - handshake ensures Core owns all memory and plan lifetimes before native - code can access them. If the adoption callback itself is interrupted, - the provider must not launch; it should make its Work terminally failed - and return normally, allowing Core to determine whether adoption took - effect before the interruption. - """ - - ... - - def submit_prevalidated( - self, - op: TransferOp, - plan_handle: object, - *, - selection: BackendIndexSelection, - notification: bytes | None, - reuse_token: object | None, - adopt_work: Callable[[BackendWork], None], - ) -> None: - """Submit a Core-authenticated immutable indexed selection. - - The default preserves provider source compatibility by forwarding the - two read-only native-int32 views to :meth:`submit`. A provider may - override this method to consume the packed views without rebuilding - them and to combine Core's proof flags with its own plan-resolution - proofs. It must not trust a duck-typed object. - """ - - if type(selection) is not BackendIndexSelection: - raise TypeError("selection must be a Core BackendIndexSelection") - selection._validate_binding(op, plan_handle) - del reuse_token - return self.submit( - op, - plan_handle, - local_indices=selection.local_indices, - remote_indices=selection.remote_indices, - notification=notification, - reuse_token=None, - adopt_work=adopt_work, - ) - - def submit_prevalidated_and_poll( - self, - op: TransferOp, - plan_handle: object, - *, - selection: BackendIndexSelection, - notification: bytes | None, - max_polls: int, - timeout_ns: int | None, - adopt_work: Callable[[BackendWork], None], - ) -> WorkState: - """Submit one authenticated selection and perform a bounded probe. - - This is capability-gated rather than composed through a callback - wrapper: only a provider that can keep preparation, post, and polling - in its intended native boundary should advertise - ``fused_prevalidated_submit_poll``. Existing providers remain source - compatible and use :meth:`submit_prevalidated` plus public Work polls. - """ - - del ( - op, - plan_handle, - selection, - notification, - max_polls, - timeout_ns, - adopt_work, - ) - raise UnsupportedError( - "backend does not support fused prevalidated submit and poll" - ) - - def create_request_slot( - self, - op: TransferOp, - plan_handle: object, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | None, - adopt_slot: Callable[..., None], - ) -> None: - """Create an opt-in persistent request for one immutable transfer. - - The provider must construct a retry-safe :class:`BackendRequestSlot` - record and synchronously call ``adopt_slot`` before its first native - mutation. A provider may pass an optional keyword-only - ``execution_dispatch`` object to that callback. Core then validates and - caches the dispatch object's execution methods while retaining the - :class:`BackendRequestSlot` wrapper for diagnostics, completion events, - and close. Omitting the keyword keeps the wrapper as the dispatch - object and is source compatible with existing providers. - - Before adoption the provider may raise only after fully rolling back - provider state. If the adoption callback is interrupted, it must not - create or launch native work and should return normally when Core may - already have retained the slot. - - A provider advertises this optional SPI with - ``capabilities.persistent_request_slots``. The default rejects it so - existing providers remain source compatible. - """ - - del op, plan_handle, local_indices, remote_indices, notification, adopt_slot - raise UnsupportedError("backend does not support persistent request slots") - - @abstractmethod - def send_notification( - self, - peer_handle: object, - payload: bytes, - *, - adopt_work: Callable[[BackendWork], None], - ) -> None: - """Send a standalone notification. - - Terminal Work completion has the meaning declared by - ``capabilities.standalone_notification_completion``. - The provider must use the same adopt-before-launch transaction described - by :meth:`submit`. - """ - - ... - - def send_notification_sync(self, peer_handle: object, payload: bytes) -> None: - """Send without creating a provider or Core Work. - - Return means the completion point declared by - ``capabilities.standalone_notification_completion``; it does not imply - remote delivery unless that declaration is ``REMOTE_DELIVERED``. - Providers opt in with ``capabilities.synchronous_notification_send``. - Under ``MULTIPLE``, Core pins the peer lifetime but does not serialize - independent native calls; providers must preserve their advertised - concurrent-entry contract. - """ - - del peer_handle, payload - raise UnsupportedError( - "backend does not support synchronous standalone notification send" - ) - - @abstractmethod - def poll_notifications( - self, max_items: int | None = None - ) -> list[BackendNotification]: - """Destructively receive up to ``max_items`` notifications. - - This v0 operation is at-most-once: the SPI has no batch token or - acknowledge/rollback phase. Providers must validate native ingress - before exposing it, while applications that require reliable delivery - must sequence, acknowledge, deduplicate, and retransmit messages. - """ - - ... - - def poll_notification_batches( - self, max_items: int | None = None - ) -> list[BackendNotificationBatch]: - """Destructively receive notifications while preserving source groups. - - ``max_items`` counts payloads, not batches. Providers with a native - grouped receive should override this method and advertise - ``capabilities.grouped_notification_receive``. This compatibility - default groups only adjacent notifications with a complete matching - identity. Anonymous notifications remain singleton batches because - scalar input cannot prove that they share a native source. - """ - - notifications = self.poll_notifications(max_items) - batches: list[BackendNotificationBatch] = [] - source: tuple[str, str] | None = None - payloads: list[bytes] = [] - - def flush() -> None: - nonlocal source, payloads - if not payloads or source is None: - return - batches.append( - BackendNotificationBatch( - payloads=tuple(payloads), - source_endpoint_id=source[0], - source_incarnation=source[1], - ) - ) - source = None - payloads = [] - - for notification in notifications: - if not isinstance(notification, BackendNotification): - raise TypeError( - "backend notifications must contain BackendNotification values" - ) - if notification.source_endpoint_id is None: - flush() - batches.append(BackendNotificationBatch((notification.payload,))) - continue - assert notification.source_incarnation is not None - next_source = ( - notification.source_endpoint_id, - notification.source_incarnation, - ) - if source != next_source: - flush() - source = next_source - payloads.append(notification.payload) - flush() - return batches - - @abstractmethod - def progress(self) -> None: ... - - @abstractmethod - def close(self) -> None: - """Close the provider; this must be idempotent after native success.""" - - ... - - -BackendFactory = Callable[ - [ - str, - str, - str, - ProgressMode, - ThreadMode, - Mapping[str, Any], - ], - Backend, -] - - -class BackendFactoryV2(Protocol): - """Factory that transfers backend ownership before native activation. - - Version-2 providers must construct a pure, closeable Python backend record, - pass it exactly once to ``adopt_backend``, and only then activate or publish - native state. A successful call returns ``None``. Core deliberately does - not inspect callable signatures or retry a factory under another contract. - """ - - def __call__( - self, - endpoint_id: str, - name: str, - incarnation: str, - progress_mode: ProgressMode, - thread_mode: ThreadMode, - options: Mapping[str, Any], - *, - adopt_backend: Callable[[Backend], None], - ) -> None: ... - - -# This is the maximum backward-compatible factory contract understood by this -# Core. A provider selecting version N may run when this value is >= N. -BACKEND_FACTORY_API_VERSION = 2 - - -@dataclass(frozen=True, slots=True) -class _BackendRegistration: - factory: BackendFactory | BackendFactoryV2 - factory_api_version: int - - -_BACKENDS: dict[str, _BackendRegistration] = {} -_BACKENDS_LOCK = threading.Lock() - - -def register_backend( - name: str, - factory: BackendFactory | BackendFactoryV2, - *, - replace: bool = False, - factory_api_version: int = 1, -) -> None: - """Register an endpoint-transfer backend factory. - - Providers should call this explicitly from an integration module. Importing - :mod:`torch` never imports an external provider. Version 1 is retained for - source compatibility and returns a :class:`Backend`; its return boundary - cannot close the final CALL-to-adoption interruption window. Version 2 - receives an ``adopt_backend`` keyword callback, must call it exactly once - before native activation, and returns ``None``. Repeating a registration - with the same factory object and exact API version is idempotent. This - closes the caller's post-commit/lost-return window without weakening name - collision protection: a different factory object or version remains busy. - """ - - if not isinstance(name, str) or not name: - raise ValueError("backend name must be a non-empty string") - if not callable(factory): - raise TypeError("backend factory must be callable") - if type(factory_api_version) is not int: - raise TypeError("factory_api_version must be an integer") - if factory_api_version not in (1, 2): - raise ValueError("factory_api_version must be 1 or 2") - with _BACKENDS_LOCK: - existing = _BACKENDS.get(name) - if existing is not None and not replace: - if ( - existing.factory is factory - and existing.factory_api_version == factory_api_version - ): - return - raise BusyError(f"transfer backend {name!r} is already registered") - _BACKENDS[name] = _BackendRegistration(factory, factory_api_version) - - -def unregister_backend(name: str) -> None: - """Remove a provider factory. - - This only affects future endpoints and is primarily intended for tests. - Existing endpoints continue to own their backend instances. - """ - - with _BACKENDS_LOCK: - if name not in _BACKENDS: - raise BackendNotFoundError(f"transfer backend {name!r} is not registered") - del _BACKENDS[name] - - -def _get_backend_registration(name: str) -> _BackendRegistration: - with _BACKENDS_LOCK: - try: - return _BACKENDS[name] - except KeyError as error: - raise BackendNotFoundError( - f"transfer backend {name!r} is not registered" - ) from error - - -def get_backend_factory(name: str) -> BackendFactory | BackendFactoryV2: - """Return the registered callable without changing the legacy query API.""" - - return _get_backend_registration(name).factory + def register(self, regions: Any, kind: str) -> Any: ... + def deregister(self, handle: Any) -> None: ... + def metadata(self) -> bytes: ... + def connect(self, metadata: bytes) -> str: ... + def disconnect(self, peer: str) -> None: ... + def prepare(self, regions: Any, kind: str, peer: str | None) -> Any: ... + def release_catalog(self, handle: Any) -> None: ... + def create_work( + self, + operation: str, + local: Any, + local_indices: Any, + remote: Any, + remote_indices: Any, + notification: bytes, + ) -> Any: ... + def post(self, handle: Any) -> str: ... + def poll(self, handle: Any) -> str: ... + def release_work(self, handle: Any) -> None: ... + def notify(self, peer: str, message: bytes) -> None: ... + def notifications(self) -> dict[str, list[bytes]]: ... diff --git a/torch/distributed/_transfer/_errors.py b/torch/distributed/_transfer/_errors.py deleted file mode 100644 index 0dd44bea6d4d6..0000000000000 --- a/torch/distributed/_transfer/_errors.py +++ /dev/null @@ -1,62 +0,0 @@ -__all__ = [ - "BackendError", - "BackendFailureError", - "BackendMismatchError", - "BackendNotFoundError", - "BusyError", - "ClosedError", - "InvalidMetadataError", - "InvalidRegionError", - "TransferCancelledError", - "TransferError", - "TransferTimeoutError", - "UnsupportedError", -] - - -class TransferError(RuntimeError): - """Base exception for endpoint transfer failures.""" - - -class BackendError(TransferError): - """Base exception for transfer backend failures.""" - - -class BackendNotFoundError(BackendError): - """Raised when a requested backend is not registered.""" - - -class BackendMismatchError(BackendError): - """Raised when metadata belongs to a different backend.""" - - -class BackendFailureError(BackendError): - """Raised when a backend operation fails.""" - - -class InvalidMetadataError(TransferError): - """Raised when endpoint metadata is malformed or unsupported.""" - - -class InvalidRegionError(TransferError): - """Raised when a memory region is invalid.""" - - -class UnsupportedError(TransferError): - """Raised when a requested operation is unsupported.""" - - -class ClosedError(TransferError): - """Raised when an operation targets a closed resource.""" - - -class BusyError(TransferError): - """Raised when a resource cannot be changed while work is in flight.""" - - -class TransferTimeoutError(TransferError): - """Raised when transfer work exceeds its deadline.""" - - -class TransferCancelledError(TransferError): - """Raised when transfer work is cancelled.""" diff --git a/torch/distributed/_transfer/_reference.py b/torch/distributed/_transfer/_reference.py index 4f898c08ff38f..70cd4e6a46571 100644 --- a/torch/distributed/_transfer/_reference.py +++ b/torch/distributed/_transfer/_reference.py @@ -1,724 +1,94 @@ -from __future__ import annotations +"""Synchronous, in-process CPU backend for examples/tests, NOT networking.""" import ctypes -import itertools -import json -import threading -from collections import deque -from dataclasses import dataclass -from typing import TYPE_CHECKING -from weakref import WeakValueDictionary - -if TYPE_CHECKING: - from collections.abc import Callable, Mapping, Sequence - from typing import Any - -from ._backend import ( - Backend, - BackendWork, - ImportedRawMemory, - ImportedRegistration, - RegisteredMemory, - register_backend, -) -from ._errors import ( - BusyError, - ClosedError, - InvalidMetadataError, - InvalidRegionError, - UnsupportedError, -) -from ._types import ( - BackendNotification, - BackendRegion, - Capabilities, - EndpointMetadata, - NotificationCompletion, - ProgressMode, - ThreadMode, - TransferOp, - WorkState, -) - - -@dataclass(eq=False, slots=True) -class _ReferencePeer: - backend: _ReferenceBackend - endpoint_id: str - incarnation: str - allowed_registration_ids: frozenset[str] - released: bool = False - - -@dataclass(eq=False, slots=True) -class _ReferencePlan: - local: tuple[BackendRegion, ...] - remote: tuple[BackendRegion, ...] - indexed: bool - released: bool = False - - -class _ReferenceWork(BackendWork): - def __init__( - self, - state: WorkState, - error: BaseException | None = None, - ) -> None: - self._state = state - self._error = error - self._closed = False - - @property - def state(self) -> WorkState: - if self._closed: - raise ClosedError("reference work is closed") - return self._state - - @property - def error(self) -> BaseException | None: - if self._closed: - raise ClosedError("reference work is closed") - return self._error - - def cancel(self) -> bool: - if self._closed: - raise ClosedError("reference work is closed") - if self._state in { - WorkState.COMPLETED, - WorkState.FAILED, - WorkState.CANCELLED, - }: - return self._state is WorkState.CANCELLED - self._state = WorkState.CANCELLED - return True - - def close(self) -> None: - self._closed = True - - -class _ReferenceBackend(Backend): - """Synchronous, in-process CPU backend used for contract testing.""" - - _CAPABILITIES = Capabilities( - read=True, - write=True, - attached_notifications=True, - standalone_notifications=True, - raw_spans=False, - strided_regions=True, - manual_progress=True, - background_progress=True, - synchronous_notification_send=True, - standalone_notification_completion=NotificationCompletion.REMOTE_DELIVERED, - ) - - def __init__( - self, - endpoint_id: str, - name: str, - incarnation: str, - progress_mode: ProgressMode, - thread_mode: ThreadMode, - options: Mapping[str, Any], - ) -> None: - self._initialize_record( - endpoint_id, - name, - incarnation, - progress_mode, - thread_mode, - options, - ) - try: - self._activate() - except BaseException as error: - try: - self.close() - except BaseException as cleanup_error: - add_note = getattr(error, "add_note", None) - if callable(add_note): - add_note( - "reference backend activation cleanup also raised: " - f"{cleanup_error}" - ) - raise - - @classmethod - def _new_unactivated( - cls, - endpoint_id: str, - name: str, - incarnation: str, - progress_mode: ProgressMode, - thread_mode: ThreadMode, - options: Mapping[str, Any], - ) -> _ReferenceBackend: - backend = cls.__new__(cls) - backend._initialize_record( - endpoint_id, - name, - incarnation, - progress_mode, - thread_mode, - options, - ) - return backend - - def _initialize_record( - self, - endpoint_id: str, - name: str, - incarnation: str, - progress_mode: ProgressMode, - thread_mode: ThreadMode, - options: Mapping[str, Any], - ) -> None: - del progress_mode, thread_mode - if options: - unknown = ", ".join(sorted(options)) - raise ValueError(f"reference backend has no options: {unknown}") - self._endpoint_id = endpoint_id - self._name = name - self._incarnation = incarnation - self._registrations: dict[str, RegisteredMemory] = {} - self._plans: set[_ReferencePlan] = set() - self._peers: set[_ReferencePeer] = set() - self._notifications: deque[BackendNotification] = deque() - self._lock = threading.RLock() - self._activation_started = False - self._activated = False - self._closed = False - - def _activate(self) -> None: - if self._closed: - raise ClosedError("cannot activate a closed reference backend") - if self._activated: - return - if self._activation_started: - raise BusyError("reference backend activation is already in progress") - self._activation_started = True - with _REFERENCE_ENDPOINTS_LOCK: - if self._endpoint_id in _REFERENCE_ENDPOINTS: - raise BusyError(f"duplicate reference endpoint {self._endpoint_id!r}") - _REFERENCE_ENDPOINTS[self._endpoint_id] = self - self._activated = True - - @property - def capabilities(self) -> Capabilities: - return self._CAPABILITIES - - @property - def supported_thread_modes(self) -> frozenset[ThreadMode]: - return frozenset( - { - ThreadMode.SINGLE, - ThreadMode.CALLER_SERIALIZED, - ThreadMode.SERIALIZED, - ThreadMode.MULTIPLE, - } - ) - - def endpoint_payload(self, registration_handles: Sequence[object]) -> bytes: - self._check_open() - registration_ids = [ - _require_registration_handle(handle) for handle in registration_handles - ] - with self._lock: - if any( - registration_id not in self._registrations - for registration_id in registration_ids - ): - raise ClosedError("reference registration is not active") - return json.dumps( - { - "endpoint_id": self._endpoint_id, - "incarnation": self._incarnation, - "registrations": sorted(registration_ids), - }, - ensure_ascii=True, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - - def register( - self, - memory: RegisteredMemory, - *, - adopt_handle: Callable[[object], None], - ) -> None: - self._check_open() - if not memory.device.startswith("cpu"): - raise UnsupportedError("reference backend supports CPU allocations only") - with self._lock: - if memory.registration_id in self._registrations: - raise BusyError("duplicate reference registration") - adopt_handle(memory.registration_id) - self._registrations[memory.registration_id] = memory - - def deregister(self, registration_handle: object) -> None: - self._check_open() - registration_id = _require_registration_handle(registration_handle) - with self._lock: - self._registrations.pop(registration_id, None) - - def describe_registration(self, registration_handle: object) -> bytes: - self._check_open() - registration_id = _require_registration_handle(registration_handle) - with self._lock: - if registration_id not in self._registrations: - raise ClosedError("reference registration is not active") - return registration_id.encode("utf-8") - - def import_peer( - self, - metadata: EndpointMetadata, - *, - adopt_handle: Callable[[object], None], - ) -> None: - self._check_open() - with _REFERENCE_ENDPOINTS_LOCK: - try: - remote = _REFERENCE_ENDPOINTS[metadata.endpoint_id] - except KeyError as error: - raise InvalidMetadataError( - "reference peer is not present in this process" - ) from error - if remote._incarnation != metadata.incarnation: - raise InvalidMetadataError("reference peer incarnation is stale") - try: - payload = json.loads(metadata.backend_payload.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as error: - raise InvalidMetadataError("reference peer payload is malformed") from error - if not isinstance(payload, dict) or set(payload) != { - "endpoint_id", - "incarnation", - "registrations", - }: - raise InvalidMetadataError("reference peer payload has invalid fields") - registration_ids = payload["registrations"] - if ( - payload["endpoint_id"] != metadata.endpoint_id - or payload["incarnation"] != metadata.incarnation - or not isinstance(registration_ids, list) - or not all(isinstance(value, str) for value in registration_ids) - or len(registration_ids) != len(set(registration_ids)) - ): - raise InvalidMetadataError("reference peer payload does not match") - advertised_ids = {item.registration_id for item in metadata.registrations} - if set(registration_ids) != advertised_ids: - raise InvalidMetadataError( - "reference payload registration grant does not match metadata" - ) - peer = _ReferencePeer( - backend=remote, - endpoint_id=metadata.endpoint_id, - incarnation=metadata.incarnation, - allowed_registration_ids=frozenset(registration_ids), - ) - adopt_handle(peer) - self._peers.add(peer) - - def release_peer(self, peer_handle: object) -> None: - self._check_open() - peer = _require_peer(peer_handle) - if peer.released: - return - peer.released = True - self._peers.discard(peer) - - def prepare( - self, - local: Sequence[BackendRegion], - remote: Sequence[BackendRegion], - *, - indexed: bool, - adopt_handle: Callable[[object], None], - ) -> None: - self._check_open() - for region in local: - if region.remote: - raise InvalidRegionError("local reference region is remote") - self._resolve_local(region) - for region in remote: - if not region.remote: - raise InvalidRegionError("remote reference region is local") - self._resolve_remote(region) - plan = _ReferencePlan(tuple(local), tuple(remote), indexed) - adopt_handle(plan) - self._plans.add(plan) - - def release_plan(self, plan_handle: object) -> None: - self._check_open() - plan = _require_plan(plan_handle) - if plan.released: - return - plan.released = True - self._plans.discard(plan) - - def submit( - self, - op: TransferOp, - plan_handle: object, - *, - local_indices: Sequence[int] | None, - remote_indices: Sequence[int] | None, - notification: bytes | None, - reuse_token: object | None, - adopt_work: Callable[[BackendWork], None], - ) -> None: - self._check_open() - plan = _require_plan(plan_handle) - if plan.released: - raise ClosedError("reference plan is released") - work = _ReferenceWork(WorkState.PENDING) - try: - adopt_work(work) - except BaseException as error: - # A callback interruption can occur immediately before or after - # Core stores the reference. Never launch in either case; if Core - # did adopt it, expose a terminal failure instead of re-raising. - work._error = error - work._state = WorkState.FAILED - return - try: - local = _select(plan.local, local_indices, indexed=plan.indexed) - remote = _select(plan.remote, remote_indices, indexed=plan.indexed) - local_spans = [ - span for region in local for span in self._resolve_local(region) - ] - remote_memory = [self._resolve_remote(region) for region in remote] - remote_spans = [span for spans, _ in remote_memory for span in spans] - if op is TransferOp.WRITE: - source, destination = local_spans, remote_spans - elif op is TransferOp.READ: - source, destination = remote_spans, local_spans - else: - raise UnsupportedError(f"unsupported transfer operation {op!r}") - _copy_spans( - source, - destination, - anchors=tuple(memory for _, memory in remote_memory), - ) - if notification is not None: - peer = self._plan_peer(plan) - peer.backend._enqueue_notification( - BackendNotification( - payload=notification, - source_endpoint_id=self._endpoint_id, - source_incarnation=self._incarnation, - ) - ) - work._state = WorkState.COMPLETED - except BaseException as error: - work._error = error - work._state = WorkState.FAILED - - def send_notification( - self, - peer_handle: object, - payload: bytes, - *, - adopt_work: Callable[[BackendWork], None], - ) -> None: - self._check_open() - work = _ReferenceWork(WorkState.PENDING) - try: - adopt_work(work) - except BaseException as error: - work._error = error - work._state = WorkState.FAILED - return - try: - peer = _require_peer(peer_handle) - if peer.released: - raise ClosedError("reference peer is released") - peer.backend._enqueue_notification( - BackendNotification( - payload=payload, - source_endpoint_id=self._endpoint_id, - source_incarnation=self._incarnation, - ) - ) - work._state = WorkState.COMPLETED - except BaseException as error: - work._error = error - work._state = WorkState.FAILED - - def send_notification_sync(self, peer_handle: object, payload: bytes) -> None: - self._check_open() - peer = _require_peer(peer_handle) - if peer.released: - raise ClosedError("reference peer is released") - peer.backend._enqueue_notification( - BackendNotification( - payload=payload, - source_endpoint_id=self._endpoint_id, - source_incarnation=self._incarnation, - ) - ) - - def poll_notifications( - self, max_items: int | None = None - ) -> list[BackendNotification]: - self._check_open() - count = len(self._notifications) if max_items is None else max_items - result: list[BackendNotification] = [] - with self._lock: - for _ in range(min(count, len(self._notifications))): - result.append(self._notifications.popleft()) - return result - - def progress(self) -> None: - self._check_open() - - def close(self) -> None: - with self._lock: - if self._closed: - return - if self._registrations or self._plans or self._peers: - raise BusyError("reference backend still owns child resources") - with _REFERENCE_ENDPOINTS_LOCK: - if _REFERENCE_ENDPOINTS.get(self._endpoint_id) is self: - _REFERENCE_ENDPOINTS.pop(self._endpoint_id, None) - self._notifications.clear() - self._activated = False - self._closed = True - - def _resolve_local(self, region: BackendRegion) -> list[tuple[int, int]]: - registration_id = _require_registration_handle(region.backend_descriptor) - with self._lock: - try: - memory = self._registrations[registration_id] - except KeyError as error: - raise ClosedError( - "local reference registration is not active" - ) from error - return _memory_spans(memory, region) - - def _resolve_remote( - self, region: BackendRegion - ) -> tuple[list[tuple[int, int]], RegisteredMemory]: - descriptor = region.backend_descriptor - if isinstance(descriptor, ImportedRawMemory): - raise UnsupportedError("reference backend rejects remote raw spans") - if not isinstance(descriptor, ImportedRegistration): - raise InvalidRegionError("invalid remote registration descriptor") - peer = _require_peer(descriptor.peer_handle) - if peer.released: - raise ClosedError("reference peer is released") - if descriptor.descriptor != descriptor.registration_id.encode("utf-8"): - raise InvalidMetadataError( - "reference registration descriptor does not match its id" - ) - if descriptor.registration_id not in peer.allowed_registration_ids: - raise InvalidMetadataError( - "remote registration is outside the imported grant" - ) - with peer.backend._lock: - try: - memory = peer.backend._registrations[descriptor.registration_id] - except KeyError as error: - raise ClosedError( - "remote reference registration is not active" - ) from error - if ( - memory.name != descriptor.name - or memory.nbytes != descriptor.nbytes - or memory.device != descriptor.device - or memory.memory_type != descriptor.memory_type - ): - raise InvalidMetadataError("remote registration attributes changed") - return _memory_spans(memory, region), memory - - def _plan_peer(self, plan: _ReferencePlan) -> _ReferencePeer: - descriptor = plan.remote[0].backend_descriptor - if not isinstance(descriptor, ImportedRegistration): - raise InvalidRegionError("reference plan has no named remote peer") - return _require_peer(descriptor.peer_handle) - - def _enqueue_notification(self, notification: BackendNotification) -> None: - with self._lock: - self._check_open() - self._notifications.append(notification) - - def _check_open(self) -> None: - if self._closed or not self._activated: - raise ClosedError(f"reference endpoint {self._name!r} is closed") - - -def _create_reference_backend( - endpoint_id: str, - name: str, - incarnation: str, - progress_mode: ProgressMode, - thread_mode: ThreadMode, - options: Mapping[str, Any], - *, - adopt_backend: Callable[[Backend], None], -) -> None: - backend = _ReferenceBackend._new_unactivated( - endpoint_id, - name, - incarnation, - progress_mode, - thread_mode, - options, - ) - adopt_backend(backend) - backend._activate() - - -def _register_reference_backend() -> None: - register_backend("reference", _create_reference_backend, factory_api_version=2) - - -def _require_registration_handle(value: object) -> str: - if not isinstance(value, str) or not value: - raise InvalidRegionError("invalid reference registration handle") - return value - - -def _require_peer(value: object) -> _ReferencePeer: - if not isinstance(value, _ReferencePeer): - raise InvalidRegionError("invalid reference peer handle") - return value - - -def _require_plan(value: object) -> _ReferencePlan: - if not isinstance(value, _ReferencePlan): - raise InvalidRegionError("invalid reference plan handle") - return value - - -def _select( - regions: tuple[BackendRegion, ...], - indices: Sequence[int] | None, - *, - indexed: bool, -) -> tuple[BackendRegion, ...]: - if indices is None: - return regions - if not indexed: - raise InvalidRegionError("indices require an indexed reference plan") - return tuple(_logical_block(regions, index) for index in indices) - - -def _logical_block( - regions: Sequence[BackendRegion], logical_index: int -) -> BackendRegion: - if isinstance(logical_index, bool) or not isinstance(logical_index, int): - raise InvalidRegionError("reference block indices must be integers") - if logical_index < 0: - raise InvalidRegionError("reference block index must be non-negative") - first_index = 0 - for region in regions: - count = region.count or 1 - if logical_index < first_index + count: - block_index = logical_index - first_index - stride = region.stride or region.nbytes - return BackendRegion( - owner=region.owner, - offset=region.offset + block_index * stride, - nbytes=region.nbytes, - device=region.device, - memory_type=region.memory_type, - backend_descriptor=region.backend_descriptor, - remote=region.remote, - ) - first_index += count - raise InvalidRegionError("reference block index is outside the region catalog") - - -def _memory_spans( - memory: RegisteredMemory, - region: BackendRegion, -) -> list[tuple[int, int]]: - count = region.count or 1 - stride = region.stride or region.nbytes - end = region.offset + (count - 1) * stride + region.nbytes - if end > memory.nbytes: - raise InvalidRegionError("reference region exceeds registered memory") - return [ - (memory.address + region.offset + index * stride, region.nbytes) - for index in range(count) - ] - - -def _copy_spans( - source: list[tuple[int, int]], - destination: list[tuple[int, int]], - *, - anchors: tuple[RegisteredMemory, ...] = (), -) -> None: - # This argument deliberately retains remote allocation owners across every - # raw pointer access, even if the exporting endpoint deregisters meanwhile. - _ = anchors - source_size = sum(nbytes for _, nbytes in source) - destination_size = sum(nbytes for _, nbytes in destination) - if source_size != destination_size: - raise InvalidRegionError( - f"source and destination sizes differ: {source_size} != {destination_size}" - ) - _reject_overlapping_spans(destination) - for source_address, source_nbytes in source: - for destination_address, destination_nbytes in destination: - if _overlaps( - source_address, - source_nbytes, - destination_address, - destination_nbytes, - ): - raise InvalidRegionError( - "reference backend rejects overlapping source and " - "destination regions" - ) - - source_index = destination_index = 0 - source_offset = destination_offset = 0 - while source_index < len(source): - source_address, source_nbytes = source[source_index] - destination_address, destination_nbytes = destination[destination_index] - chunk = min( - source_nbytes - source_offset, - destination_nbytes - destination_offset, - ) - ctypes.memmove( - destination_address + destination_offset, - source_address + source_offset, - chunk, - ) - source_offset += chunk - destination_offset += chunk - if source_offset == source_nbytes: - source_index += 1 - source_offset = 0 - if destination_offset == destination_nbytes: - destination_index += 1 - destination_offset = 0 - - -def _reject_overlapping_spans(spans: list[tuple[int, int]]) -> None: - intervals = sorted((address, address + nbytes) for address, nbytes in spans) - for previous, current in itertools.pairwise(intervals): - if current[0] < previous[1]: - raise InvalidRegionError( - "reference backend rejects overlapping destination regions" - ) - - -def _overlaps( - first_address: int, - first_nbytes: int, - second_address: int, - second_nbytes: int, -) -> bool: - return ( - first_address < second_address + second_nbytes - and second_address < first_address + first_nbytes - ) - - -_REFERENCE_ENDPOINTS: WeakValueDictionary[str, _ReferenceBackend] = ( - WeakValueDictionary() -) -_REFERENCE_ENDPOINTS_LOCK = threading.Lock() +import uuid +import weakref + + +class MemoryBackend: + _agents = weakref.WeakValueDictionary() + + def __init__(self): + self.name = uuid.uuid4().hex + self._agents[self.name] = self + self.regions = {} + self.messages = {} + + def register(self, regions, kind): + if kind != "host": + raise ValueError("reference backend supports CPU memory only") + handle = object() + self.regions[handle] = tuple(regions) + return handle + + def deregister(self, handle): + del self.regions[handle] + + def metadata(self): + return self.name.encode() + + def connect(self, metadata): + name = metadata.decode() + self._agents[name] # fail immediately if the peer is absent + return name + + def disconnect(self, peer): + pass + + def prepare(self, regions, kind, peer): + if kind != "host": + raise ValueError("reference backend supports CPU memory only") + owner = self if peer is None else self._agents[peer] + expanded = [] + for row in regions: + address, length, device = map(int, row[:3]) + stride, count = map(int, row[3:]) if len(row) == 5 else (length, 1) + if length <= 0 or count <= 0 or stride < length: + raise ValueError("invalid region") + for i in range(count): + span = (address + i * stride, length, device) + if not any( + a <= span[0] and span[0] + length <= a + n and d == device + for rows in owner.regions.values() + for a, n, d in rows + ): + raise ValueError("descriptor outside registered memory") + expanded.append(span) + return owner, expanded + + def release_catalog(self, handle): + pass + + def create_work( + self, operation, local, local_indices, remote, remote_indices, notification + ): + pairs = [] + for li, ri in zip(local_indices, remote_indices): + if not (0 <= li < len(local[1]) and 0 <= ri < len(remote[1])): + raise ValueError("index outside catalog") + left, right = local[1][li], remote[1][ri] + if left[1] != right[1]: + raise ValueError("descriptor lengths differ") + pairs.append((left, right) if operation == "WRITE" else (right, left)) + return pairs, remote[0], notification + + def post(self, handle): + pairs, peer, message = handle + for source, target in pairs: + ctypes.memmove(target[0], source[0], source[1]) + if message: + peer.messages.setdefault(self.name, []).append(message) + return "done" + + def poll(self, handle): + return "done" + + def release_work(self, handle): + pass + + def notify(self, peer, message): + self._agents[peer].messages.setdefault(self.name, []).append(message) + + def notifications(self): + messages, self.messages = self.messages, {} + return messages diff --git a/torch/distributed/_transfer/_types.py b/torch/distributed/_transfer/_types.py deleted file mode 100644 index f696d01f1db67..0000000000000 --- a/torch/distributed/_transfer/_types.py +++ /dev/null @@ -1,867 +0,0 @@ -import base64 -import binascii -import json -from collections.abc import Mapping -from dataclasses import dataclass, field -from enum import Enum -from typing import Any, ClassVar - -from ._errors import InvalidMetadataError, InvalidRegionError - - -__all__ = [ - "BackendNotification", - "BackendNotificationBatch", - "BackendRegion", - "Capabilities", - "CudaOrderingMode", - "EndpointIdentity", - "EndpointInfo", - "EndpointMetadata", - "Notification", - "NotificationBatch", - "NotificationCompletion", - "ProgressMode", - "RegistrationMetadata", - "ThreadMode", - "TransferOp", - "WorkState", -] - - -class WorkState(str, Enum): - PENDING = "pending" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class TransferOp(str, Enum): - READ = "read" - WRITE = "write" - - -class ProgressMode(str, Enum): - MANUAL = "manual" - BACKGROUND = "background" - - -class ThreadMode(str, Enum): - """Concurrency contract for one Endpoint resource graph. - - ``SINGLE`` is confined to the creating thread. ``CALLER_SERIALIZED`` may - migrate between threads, but the caller must prevent overlapping or - reentrant operations across the endpoint and every derived resource. - ``SERIALIZED`` asks Core to provide that exclusion. ``MULTIPLE`` permits - concurrent calls and requires a provider that supports them. - """ - - SINGLE = "single" - CALLER_SERIALIZED = "caller_serialized" - SERIALIZED = "serialized" - MULTIPLE = "multiple" - - -class CudaOrderingMode(str, Enum): - """How transport access is ordered after CUDA producer work.""" - - CALLER_READY = "caller_ready" - HOST_WAIT = "host_wait" - - -class NotificationCompletion(str, Enum): - """Meaning of terminal standalone-notification Work.""" - - LOCAL_ACCEPTED = "local_accepted" - REMOTE_DELIVERED = "remote_delivered" - - -@dataclass(frozen=True, slots=True) -class EndpointIdentity: - """Stable logical endpoint plus one process-lifetime incarnation.""" - - endpoint_id: str - incarnation: str - - def __post_init__(self) -> None: - _require_nonempty_str(self.endpoint_id, "endpoint_id") - _require_nonempty_str(self.incarnation, "incarnation") - - -@dataclass(frozen=True, slots=True) -class Capabilities: - read: bool = False - write: bool = False - attached_notifications: bool = False - standalone_notifications: bool = False - raw_spans: bool = False - strided_regions: bool = False - manual_progress: bool = False - background_progress: bool = False - cuda_completion_event: bool = False - persistent_request_slots: bool = False - fused_request_slot_start_poll: bool = False - fused_request_slot_poll: bool = False - cuda_ordering_modes: frozenset[CudaOrderingMode] = frozenset( - {CudaOrderingMode.CALLER_READY} - ) - standalone_notification_completion: NotificationCompletion | None = None - deferred_successful_request_slot_recycle: bool = False - # Appended to preserve positional construction compatibility for every - # pre-existing capability field. New code should use keyword arguments. - request_slot_notification_overrides: bool = False - fused_prevalidated_submit_poll: bool = False - fused_work_poll: bool = False - grouped_notification_receive: bool = False - synchronous_notification_send: bool = False - - def __post_init__(self) -> None: - values = { - "read": self.read, - "write": self.write, - "attached_notifications": self.attached_notifications, - "standalone_notifications": self.standalone_notifications, - "raw_spans": self.raw_spans, - "strided_regions": self.strided_regions, - "manual_progress": self.manual_progress, - "background_progress": self.background_progress, - "cuda_completion_event": self.cuda_completion_event, - "persistent_request_slots": self.persistent_request_slots, - "request_slot_notification_overrides": ( - self.request_slot_notification_overrides - ), - "fused_prevalidated_submit_poll": (self.fused_prevalidated_submit_poll), - "fused_work_poll": self.fused_work_poll, - "grouped_notification_receive": self.grouped_notification_receive, - "synchronous_notification_send": self.synchronous_notification_send, - "fused_request_slot_start_poll": self.fused_request_slot_start_poll, - "fused_request_slot_poll": self.fused_request_slot_poll, - "deferred_successful_request_slot_recycle": ( - self.deferred_successful_request_slot_recycle - ), - } - for name, value in values.items(): - if not isinstance(value, bool): - raise TypeError(f"capability {name} must be a bool") - try: - ordering_modes = frozenset(self.cuda_ordering_modes) - except TypeError as error: - raise TypeError("cuda_ordering_modes must be an iterable") from error - if not ordering_modes or not all( - isinstance(mode, CudaOrderingMode) for mode in ordering_modes - ): - raise TypeError("cuda_ordering_modes must contain CudaOrderingMode values") - if CudaOrderingMode.CALLER_READY not in ordering_modes: - raise ValueError("cuda_ordering_modes must include caller_ready") - object.__setattr__(self, "cuda_ordering_modes", ordering_modes) - if self.fused_request_slot_start_poll and not self.persistent_request_slots: - raise ValueError( - "fused_request_slot_start_poll requires persistent_request_slots" - ) - if ( - self.request_slot_notification_overrides - and not self.persistent_request_slots - ): - raise ValueError( - "request_slot_notification_overrides requires " - "persistent_request_slots" - ) - if self.request_slot_notification_overrides and not self.attached_notifications: - raise ValueError( - "request_slot_notification_overrides requires " "attached_notifications" - ) - if self.fused_request_slot_poll and not self.persistent_request_slots: - raise ValueError( - "fused_request_slot_poll requires persistent_request_slots" - ) - if ( - self.deferred_successful_request_slot_recycle - and not self.persistent_request_slots - ): - raise ValueError( - "deferred_successful_request_slot_recycle requires " - "persistent_request_slots" - ) - completion = self.standalone_notification_completion - if self.standalone_notifications and not isinstance( - completion, NotificationCompletion - ): - raise TypeError( - "notification-capable backends must declare standalone " - "notification completion" - ) - if not self.standalone_notifications and completion is not None: - raise ValueError( - "standalone notification completion requires notifications" - ) - if self.grouped_notification_receive and not self.notifications: - raise ValueError("grouped_notification_receive requires notifications") - if self.synchronous_notification_send and not self.standalone_notifications: - raise ValueError( - "synchronous_notification_send requires standalone_notifications" - ) - - @property - def notifications(self) -> bool: - """Whether either attached or standalone notifications are available.""" - - return self.attached_notifications or self.standalone_notifications - - -@dataclass(frozen=True, slots=True) -class EndpointInfo: - backend: str - endpoint_id: str - name: str - incarnation: str - capabilities: Capabilities - - def __post_init__(self) -> None: - _require_nonempty_str(self.backend, "backend") - _require_nonempty_str(self.endpoint_id, "endpoint_id") - _require_nonempty_str(self.name, "name") - _require_nonempty_str(self.incarnation, "incarnation") - if not isinstance(self.capabilities, Capabilities): - raise TypeError("capabilities must be Capabilities") - - @property - def identity(self) -> EndpointIdentity: - return EndpointIdentity(self.endpoint_id, self.incarnation) - - -@dataclass(frozen=True, slots=True) -class BackendRegion: - """Provider-facing region with a core-owned lifetime anchor. - - Providers consume ``backend_descriptor`` and must treat ``owner`` as opaque. - For local regions the descriptor is the provider registration handle; for - remote regions it is ``ImportedRegistration`` or ``ImportedRawMemory``. - In an indexed plan, ``count`` compactly contributes that many independently - selectable logical blocks, ordered by increasing stride displacement. - """ - - owner: object = field(repr=False) - offset: int - nbytes: int - device: str - memory_type: str - backend_descriptor: object = field(repr=False) - remote: bool = False - stride: int | None = None - count: int | None = None - - def __post_init__(self) -> None: - if not _is_int(self.offset) or self.offset < 0: - raise InvalidRegionError("offset must be a non-negative integer") - if not _is_int(self.nbytes) or self.nbytes <= 0: - raise InvalidRegionError("nbytes must be a positive integer") - _require_nonempty_str(self.device, "device", InvalidRegionError) - _require_nonempty_str(self.memory_type, "memory_type", InvalidRegionError) - if not isinstance(self.remote, bool): - raise InvalidRegionError("remote must be a bool") - if (self.stride is None) != (self.count is None): - raise InvalidRegionError("stride and count must be specified together") - if self.stride is not None: - if not _is_int(self.stride) or self.stride <= 0: - raise InvalidRegionError("stride must be a positive integer") - if ( - isinstance(self.count, bool) - or not isinstance(self.count, int) - or self.count <= 0 - ): - raise InvalidRegionError("count must be a positive integer") - if self.stride < self.nbytes: - raise InvalidRegionError("stride must be at least nbytes") - - -@dataclass(frozen=True, slots=True) -class BackendNotification: - MAX_PAYLOAD_BYTES: ClassVar[int] = 1024 * 1024 - - payload: bytes = field(repr=False) - source_endpoint_id: str | None = None - source_incarnation: str | None = None - - def __post_init__(self) -> None: - payload = _immutable_bytes(self.payload, "payload") - if len(payload) > self.MAX_PAYLOAD_BYTES: - raise ValueError("notification payload exceeds the wire size limit") - object.__setattr__(self, "payload", payload) - if self.source_endpoint_id is not None: - _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") - if self.source_incarnation is not None: - _require_nonempty_str(self.source_incarnation, "source_incarnation") - if (self.source_endpoint_id is None) != (self.source_incarnation is None): - raise ValueError( - "notification source_endpoint_id and source_incarnation must be " - "specified together" - ) - - @property - def source(self) -> EndpointIdentity | None: - if self.source_endpoint_id is None or self.source_incarnation is None: - return None - return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) - - -@dataclass(frozen=True, slots=True) -class BackendNotificationBatch: - """Provider-facing notification group from one native source batch. - - Providers may return multiple batches for one source. Payload order is - significant within and across that source's batches; no total order is - promised between different sources. - """ - - MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES - - payloads: tuple[bytes, ...] = field(repr=False) - total_nbytes: int = field(init=False) - source_endpoint_id: str | None = None - source_incarnation: str | None = None - - def __post_init__(self) -> None: - payloads, total_nbytes = _immutable_notification_payloads( - self.payloads, self.MAX_PAYLOAD_BYTES - ) - if not payloads: - raise ValueError("notification batch must contain at least one payload") - object.__setattr__(self, "payloads", payloads) - object.__setattr__(self, "total_nbytes", total_nbytes) - if self.source_endpoint_id is not None: - _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") - if self.source_incarnation is not None: - _require_nonempty_str(self.source_incarnation, "source_incarnation") - if (self.source_endpoint_id is None) != (self.source_incarnation is None): - raise ValueError( - "notification source_endpoint_id and source_incarnation must be " - "specified together" - ) - - @property - def source(self) -> EndpointIdentity | None: - if self.source_endpoint_id is None or self.source_incarnation is None: - return None - return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) - - -@dataclass(frozen=True, slots=True) -class Notification: - MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES - - payload: bytes = field(repr=False) - source_endpoint_id: str | None = None - source_incarnation: str | None = None - - def __post_init__(self) -> None: - payload = _immutable_bytes(self.payload, "payload") - if len(payload) > self.MAX_PAYLOAD_BYTES: - raise ValueError("notification payload exceeds the wire size limit") - object.__setattr__(self, "payload", payload) - if self.source_endpoint_id is not None: - _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") - if self.source_incarnation is not None: - _require_nonempty_str(self.source_incarnation, "source_incarnation") - if (self.source_endpoint_id is None) != (self.source_incarnation is None): - raise ValueError( - "notification source_endpoint_id and source_incarnation must be " - "specified together" - ) - - @property - def source(self) -> EndpointIdentity | None: - if self.source_endpoint_id is None or self.source_incarnation is None: - return None - return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) - - -@dataclass(frozen=True, slots=True) -class NotificationBatch: - """Public notification group preserving one provider source boundary.""" - - MAX_PAYLOAD_BYTES: ClassVar[int] = BackendNotification.MAX_PAYLOAD_BYTES - - payloads: tuple[bytes, ...] = field(repr=False) - source_endpoint_id: str | None = None - source_incarnation: str | None = None - - def __post_init__(self) -> None: - payloads, _ = _immutable_notification_payloads( - self.payloads, self.MAX_PAYLOAD_BYTES - ) - if not payloads: - raise ValueError("notification batch must contain at least one payload") - object.__setattr__(self, "payloads", payloads) - if self.source_endpoint_id is not None: - _require_nonempty_str(self.source_endpoint_id, "source_endpoint_id") - if self.source_incarnation is not None: - _require_nonempty_str(self.source_incarnation, "source_incarnation") - if (self.source_endpoint_id is None) != (self.source_incarnation is None): - raise ValueError( - "notification source_endpoint_id and source_incarnation must be " - "specified together" - ) - - @property - def source(self) -> EndpointIdentity | None: - if self.source_endpoint_id is None or self.source_incarnation is None: - return None - return EndpointIdentity(self.source_endpoint_id, self.source_incarnation) - - -def _notification_batch_from_backend( - batch: BackendNotificationBatch, -) -> NotificationBatch: - """Adopt one exact, already-validated provider batch without rescanning it.""" - - if type(batch) is not BackendNotificationBatch: - raise TypeError("expected an exact BackendNotificationBatch") - result = object.__new__(NotificationBatch) - object.__setattr__(result, "payloads", batch.payloads) - object.__setattr__(result, "source_endpoint_id", batch.source_endpoint_id) - object.__setattr__(result, "source_incarnation", batch.source_incarnation) - return result - - -@dataclass(frozen=True, slots=True) -class RegistrationMetadata: - registration_id: str - name: str - nbytes: int - device: str - memory_type: str - backend_descriptor: bytes = field(repr=False) - - def __post_init__(self) -> None: - _require_nonempty_str( - self.registration_id, "registration_id", InvalidMetadataError - ) - _require_nonempty_str(self.name, "name", InvalidMetadataError) - if not _is_int(self.nbytes) or self.nbytes <= 0: - raise InvalidMetadataError("registration nbytes must be positive") - _require_nonempty_str(self.device, "device", InvalidMetadataError) - _require_nonempty_str(self.memory_type, "memory_type", InvalidMetadataError) - descriptor = _immutable_bytes( - self.backend_descriptor, "backend_descriptor", InvalidMetadataError - ) - object.__setattr__(self, "backend_descriptor", descriptor) - - -@dataclass(frozen=True, slots=True) -class EndpointMetadata: - WIRE_MAJOR: ClassVar[int] = 1 - WIRE_MINOR: ClassVar[int] = 0 - KNOWN_REQUIRED_FEATURES: ClassVar[frozenset[str]] = frozenset() - MAX_ENCODED_BYTES: ClassVar[int] = 64 * 1024 * 1024 - MAX_BACKEND_PAYLOAD_BYTES: ClassVar[int] = 32 * 1024 * 1024 - MAX_DESCRIPTOR_BYTES: ClassVar[int] = 4 * 1024 * 1024 - MAX_TOTAL_DESCRIPTOR_BYTES: ClassVar[int] = 32 * 1024 * 1024 - MAX_REGISTRATIONS: ClassVar[int] = 65_536 - MAX_REQUIRED_FEATURES: ClassVar[int] = 256 - MAX_EXTENSIONS: ClassVar[int] = 256 - MAX_EXTENSION_BYTES: ClassVar[int] = 4 * 1024 * 1024 - MAX_TOTAL_EXTENSION_BYTES: ClassVar[int] = 32 * 1024 * 1024 - MAX_STRING_BYTES: ClassVar[int] = 4_096 - - wire_major: int - wire_minor: int - required_features: tuple[str, ...] - backend: str - endpoint_id: str - name: str - incarnation: str - backend_payload: bytes = field(repr=False) - registrations: tuple[RegistrationMetadata, ...] - extensions: tuple[tuple[str, bytes], ...] = field(default=(), repr=False) - - def __post_init__(self) -> None: - if not _is_int(self.wire_major) or self.wire_major != self.WIRE_MAJOR: - raise InvalidMetadataError( - f"unsupported endpoint metadata wire major: {self.wire_major!r}" - ) - if not _is_int(self.wire_minor) or self.wire_minor < 0: - raise InvalidMetadataError( - f"invalid endpoint metadata wire minor: {self.wire_minor!r}" - ) - try: - required_features = tuple(self.required_features) - except TypeError as error: - raise InvalidMetadataError( - "required_features must be an iterable" - ) from error - if not all(isinstance(item, str) and item for item in required_features): - raise InvalidMetadataError( - "required_features must contain non-empty strings" - ) - if len(required_features) > self.MAX_REQUIRED_FEATURES: - raise InvalidMetadataError("too many required endpoint metadata features") - if any( - _utf8_nbytes(item, "required feature") > self.MAX_STRING_BYTES - for item in required_features - ): - raise InvalidMetadataError( - "required endpoint metadata feature exceeds the wire size limit" - ) - if len(required_features) != len(set(required_features)): - raise InvalidMetadataError("required_features must be unique") - unknown_features = set(required_features) - self.KNOWN_REQUIRED_FEATURES - if unknown_features: - raise InvalidMetadataError( - "unsupported required endpoint metadata features: " - f"{', '.join(sorted(unknown_features))}" - ) - object.__setattr__(self, "required_features", tuple(sorted(required_features))) - _require_nonempty_str(self.backend, "backend", InvalidMetadataError) - _require_nonempty_str(self.endpoint_id, "endpoint_id", InvalidMetadataError) - _require_nonempty_str(self.name, "name", InvalidMetadataError) - _require_nonempty_str(self.incarnation, "incarnation", InvalidMetadataError) - for field_name in ("backend", "endpoint_id", "name", "incarnation"): - if ( - _utf8_nbytes(getattr(self, field_name), field_name) - > self.MAX_STRING_BYTES - ): - raise InvalidMetadataError(f"{field_name} exceeds the wire size limit") - payload = _immutable_bytes( - self.backend_payload, "backend_payload", InvalidMetadataError - ) - if len(payload) > self.MAX_BACKEND_PAYLOAD_BYTES: - raise InvalidMetadataError("backend_payload exceeds the wire size limit") - object.__setattr__(self, "backend_payload", payload) - - try: - registrations = tuple(self.registrations) - except TypeError as error: - raise InvalidMetadataError("registrations must be an iterable") from error - if not all(isinstance(item, RegistrationMetadata) for item in registrations): - raise InvalidMetadataError( - "registrations must contain RegistrationMetadata values" - ) - registration_ids = [item.registration_id for item in registrations] - names = [item.name for item in registrations] - if len(registration_ids) != len(set(registration_ids)): - raise InvalidMetadataError("registration_id values must be unique") - if len(names) != len(set(names)): - raise InvalidMetadataError("registration names must be unique") - if len(registrations) > self.MAX_REGISTRATIONS: - raise InvalidMetadataError("too many endpoint metadata registrations") - descriptor_bytes = 0 - for registration in registrations: - for field_name in ( - "registration_id", - "name", - "device", - "memory_type", - ): - if ( - _utf8_nbytes( - getattr(registration, field_name), - f"registration {field_name}", - ) - > self.MAX_STRING_BYTES - ): - raise InvalidMetadataError( - f"registration {field_name} exceeds the wire size limit" - ) - if len(registration.backend_descriptor) > self.MAX_DESCRIPTOR_BYTES: - raise InvalidMetadataError( - "registration descriptor exceeds the wire size limit" - ) - descriptor_bytes += len(registration.backend_descriptor) - if descriptor_bytes > self.MAX_TOTAL_DESCRIPTOR_BYTES: - raise InvalidMetadataError( - "registration descriptors exceed the aggregate wire size limit" - ) - object.__setattr__(self, "registrations", registrations) - - raw_extensions: Any = self.extensions - if isinstance(raw_extensions, Mapping): - extension_items = tuple(raw_extensions.items()) - else: - try: - extension_items = tuple(raw_extensions) - except TypeError as error: - raise InvalidMetadataError("extensions must be a mapping") from error - if len(extension_items) > self.MAX_EXTENSIONS: - raise InvalidMetadataError("too many endpoint metadata extensions") - extensions: list[tuple[str, bytes]] = [] - extension_bytes = 0 - for item in extension_items: - if not isinstance(item, tuple) or len(item) != 2: - raise InvalidMetadataError("extensions must contain key/value pairs") - key, value = item - if not isinstance(key, str) or not key: - raise InvalidMetadataError("extension names must be non-empty strings") - if _utf8_nbytes(key, "extension name") > self.MAX_STRING_BYTES: - raise InvalidMetadataError("extension name exceeds the wire size limit") - extension = _immutable_bytes( - value, f"extension {key!r}", InvalidMetadataError - ) - if len(extension) > self.MAX_EXTENSION_BYTES: - raise InvalidMetadataError( - f"extension {key!r} exceeds the wire size limit" - ) - extension_bytes += len(extension) - extensions.append((key, extension)) - if extension_bytes > self.MAX_TOTAL_EXTENSION_BYTES: - raise InvalidMetadataError( - "endpoint metadata extensions exceed the aggregate wire size limit" - ) - extension_names = [key for key, _ in extensions] - if len(extension_names) != len(set(extension_names)): - raise InvalidMetadataError("extension names must be unique") - object.__setattr__(self, "extensions", tuple(sorted(extensions))) - - def __bytes__(self) -> bytes: - registrations = [ - { - "backend_descriptor": _encode_bytes(item.backend_descriptor), - "device": item.device, - "memory_type": item.memory_type, - "name": item.name, - "nbytes": item.nbytes, - "registration_id": item.registration_id, - } - for item in self.registrations - ] - value = { - "backend": self.backend, - "backend_payload": _encode_bytes(self.backend_payload), - "endpoint_id": self.endpoint_id, - "extensions": { - key: _encode_bytes(extension) for key, extension in self.extensions - }, - "incarnation": self.incarnation, - "name": self.name, - "registrations": registrations, - "required_features": list(self.required_features), - "wire_version": { - "major": self.wire_major, - "minor": self.wire_minor, - }, - } - encoded = json.dumps( - value, ensure_ascii=True, separators=(",", ":"), sort_keys=True - ).encode("utf-8") - if len(encoded) > self.MAX_ENCODED_BYTES: - raise InvalidMetadataError("endpoint metadata exceeds the wire size limit") - return encoded - - def to_bytes(self) -> bytes: - """Return the canonical wire representation.""" - - return bytes(self) - - @classmethod - def from_bytes(cls, data: bytes | bytearray | memoryview) -> "EndpointMetadata": - if not isinstance(data, (bytes, bytearray, memoryview)): - raise InvalidMetadataError("endpoint metadata must be bytes-like") - try: - encoded_nbytes = data.nbytes if isinstance(data, memoryview) else len(data) - except ValueError as error: - raise InvalidMetadataError( - "endpoint metadata buffer is unavailable" - ) from error - if encoded_nbytes > cls.MAX_ENCODED_BYTES: - raise InvalidMetadataError("endpoint metadata exceeds the wire size limit") - try: - encoded = bytes(data) - if len(encoded) > cls.MAX_ENCODED_BYTES: - raise InvalidMetadataError( - "endpoint metadata exceeds the wire size limit" - ) - value = json.loads(encoded.decode("utf-8"), object_pairs_hook=_json_object) - except ( - UnicodeDecodeError, - ValueError, - RecursionError, - InvalidMetadataError, - ) as error: - if isinstance(error, InvalidMetadataError): - raise - raise InvalidMetadataError("endpoint metadata is not valid JSON") from error - - _require_object_keys( - value, - { - "backend", - "backend_payload", - "endpoint_id", - "extensions", - "incarnation", - "name", - "registrations", - "required_features", - "wire_version", - }, - "endpoint metadata", - ) - wire_version = value["wire_version"] - _require_object_keys(wire_version, {"major", "minor"}, "wire_version") - required_features = value["required_features"] - if not isinstance(required_features, list): - raise InvalidMetadataError("required_features must be a JSON array") - if len(required_features) > cls.MAX_REQUIRED_FEATURES: - raise InvalidMetadataError("too many required endpoint metadata features") - raw_extensions = value["extensions"] - if not isinstance(raw_extensions, dict): - raise InvalidMetadataError("extensions must be a JSON object") - if len(raw_extensions) > cls.MAX_EXTENSIONS: - raise InvalidMetadataError("too many endpoint metadata extensions") - raw_registrations = value["registrations"] - if not isinstance(raw_registrations, list): - raise InvalidMetadataError("registrations must be a JSON array") - if len(raw_registrations) > cls.MAX_REGISTRATIONS: - raise InvalidMetadataError("too many endpoint metadata registrations") - registrations = tuple( - _registration_from_json(item, index) - for index, item in enumerate(raw_registrations) - ) - return cls( - wire_major=wire_version["major"], - wire_minor=wire_version["minor"], - required_features=tuple(required_features), - backend=value["backend"], - endpoint_id=value["endpoint_id"], - name=value["name"], - incarnation=value["incarnation"], - backend_payload=_decode_bytes(value["backend_payload"], "backend_payload"), - registrations=registrations, - extensions=tuple( - (key, _decode_bytes(extension, f"extension {key!r}")) - for key, extension in raw_extensions.items() - ), - ) - - -def _registration_from_json(value: Any, index: int) -> RegistrationMetadata: - context = f"registration at index {index}" - _require_object_keys( - value, - { - "backend_descriptor", - "device", - "memory_type", - "name", - "nbytes", - "registration_id", - }, - context, - ) - return RegistrationMetadata( - registration_id=value["registration_id"], - name=value["name"], - nbytes=value["nbytes"], - device=value["device"], - memory_type=value["memory_type"], - backend_descriptor=_decode_bytes( - value["backend_descriptor"], f"{context} backend_descriptor" - ), - ) - - -def _json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value: dict[str, Any] = {} - for key, item in pairs: - if key in value: - raise InvalidMetadataError(f"duplicate JSON field: {key!r}") - value[key] = item - return value - - -def _require_object_keys(value: Any, expected: set[str], context: str) -> None: - if not isinstance(value, dict): - raise InvalidMetadataError(f"{context} must be a JSON object") - keys = set(value) - missing = expected - keys - unknown = keys - expected - if missing: - raise InvalidMetadataError( - f"{context} is missing fields: {', '.join(sorted(missing))}" - ) - if unknown: - raise InvalidMetadataError( - f"{context} has unknown fields: {', '.join(sorted(unknown))}" - ) - - -def _encode_bytes(value: bytes) -> str: - return base64.b64encode(value).decode("ascii") - - -def _decode_bytes(value: Any, field: str) -> bytes: - if not isinstance(value, str): - raise InvalidMetadataError(f"{field} must be a base64 string") - try: - encoded = value.encode("ascii") - return base64.b64decode(encoded, validate=True) - except (UnicodeEncodeError, binascii.Error, ValueError) as error: - raise InvalidMetadataError(f"{field} is not valid base64") from error - - -def _immutable_bytes( - value: Any, - field: str, - error_type: type[Exception] = TypeError, -) -> bytes: - if not isinstance(value, (bytes, bytearray, memoryview)): - raise error_type(f"{field} must be bytes-like") - return bytes(value) - - -def _immutable_notification_payloads( - value: Any, - max_payload_bytes: int, -) -> tuple[tuple[bytes, ...], int]: - """Snapshot payloads and count bytes during the required size validation.""" - - try: - payloads = value if type(value) is tuple else tuple(value) - except TypeError as error: - raise TypeError("payloads must be an iterable of bytes-like values") from error - total_nbytes = 0 - for payload in payloads: - if type(payload) is not bytes: - for item in payloads: - if not isinstance(item, (bytes, bytearray, memoryview)): - raise TypeError("payload must be bytes-like") - reported_nbytes = ( - item.nbytes if isinstance(item, memoryview) else len(item) - ) - if reported_nbytes > max_payload_bytes: - raise ValueError("notification payload exceeds the wire size limit") - immutable_payloads = tuple( - _immutable_bytes(item, "payload") for item in payloads - ) - total_nbytes = 0 - for item in immutable_payloads: - item_nbytes = len(item) - if item_nbytes > max_payload_bytes: - raise ValueError("notification payload exceeds the wire size limit") - total_nbytes += item_nbytes - return immutable_payloads, total_nbytes - payload_nbytes = len(payload) - if payload_nbytes > max_payload_bytes: - raise ValueError("notification payload exceeds the wire size limit") - total_nbytes += payload_nbytes - return payloads, total_nbytes - - -def _require_nonempty_str( - value: Any, - field: str, - error_type: type[Exception] = TypeError, -) -> None: - if not isinstance(value, str) or not value: - raise error_type(f"{field} must be a non-empty string") - - -def _utf8_nbytes(value: str, field: str) -> int: - try: - return len(value.encode("utf-8")) - except UnicodeEncodeError as error: - raise InvalidMetadataError(f"{field} is not valid Unicode") from error - - -def _is_int(value: Any) -> bool: - return isinstance(value, int) and not isinstance(value, bool)