diff --git a/test/distributed/test_transfer.py b/test/distributed/test_transfer.py new file mode 100644 index 0000000000000..1d75962e9acdc --- /dev/null +++ b/test/distributed/test_transfer.py @@ -0,0 +1,127 @@ +import gc +import unittest +import weakref +from unittest.mock import Mock + +import torch +from torch.distributed._transfer import Endpoint +from torch.distributed._transfer._reference import MemoryBackend + + +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 tearDown(self): + self.a.close() + self.b.close() + + def submit(self, operation="WRITE", **kwargs): + return self.a.submit( + operation, self.local, [1, 0], self.remote, [0, 1], **kwargs + ) + + 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() + + 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.backend.poll.return_value = "done" + self.backend.release_work = Mock(side_effect=RuntimeError("retry")) + with self.assertRaises(RuntimeError): + work.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_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) + + def test_invalid_catalogs_indices_and_registration(self): + with self.assertRaises(ValueError): + self.a.prepare([(self.src.data_ptr() + 8, 4, 0)], kind="host") + with self.assertRaises(ValueError): + self.a.submit("WRITE", self.remote, [0], self.local, [0]) + with self.assertRaises(ValueError): + 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__": + unittest.main() diff --git a/torch/distributed/_transfer/README.md b/torch/distributed/_transfer/README.md new file mode 100644 index 0000000000000..3b2392732c07d --- /dev/null +++ b/torch/distributed/_transfer/README.md @@ -0,0 +1,54 @@ +# Transfer API — review MVP + +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 +import torch +from torch.distributed._transfer import Endpoint +from torch.distributed._transfer._reference import MemoryBackend + +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() +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 new file mode 100644 index 0000000000000..d5e117a599eaa --- /dev/null +++ b/torch/distributed/_transfer/__init__.py @@ -0,0 +1,9 @@ +"""Experimental, backend-independent point-to-point memory transfers. + +This is a review prototype, not a stable or production-supported API. +""" + +from ._api import Catalog, Endpoint, Registration, Work +from ._backend import Backend + +__all__ = ["Backend", "Catalog", "Endpoint", "Registration", "Work"] diff --git a/torch/distributed/_transfer/_api.py b/torch/distributed/_transfer/_api.py new file mode 100644 index 0000000000000..1e10d2688f07e --- /dev/null +++ b/torch/distributed/_transfer/_api.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import threading +import time +from typing import Any + +import torch + +from ._backend import 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() + + +class Endpoint: + """One provider instance and its locally owned resources. + + 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: Backend): + self._backend = backend + self._lock = threading.RLock() + self._registrations: set[Registration] = set() + self._catalogs: set[Catalog] = set() + self._works: set[Work] = set() + self._peers: set[str] = set() + self._closed = False + + def _check_open(self): + if self._closed: + raise RuntimeError("endpoint is closed") + + 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: + 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 metadata(self) -> bytes: + """Export provider metadata AFTER registration; exchange out of band.""" + with self._lock: + self._check_open() + return self._backend.metadata() + + def connect(self, metadata: bytes) -> str: + with self._lock: + self._check_open() + peer = self._backend.connect(metadata) + self._peers.add(peer) + return peer + + def disconnect(self, peer: str) -> None: + with self._lock: + 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 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: + 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 submit( + self, + operation: str, + local: Catalog, + local_indices: Any, + remote: Catalog, + remote_indices: Any, + *, + notification: bytes = b"", + ordering: str = "caller_ready", + ready_event: Any = None, + ) -> Work: + """Submit READ (remote -> local) or WRITE (local -> remote). + + 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, + ) + 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 + + 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 notifications(self) -> dict[str, list[bytes]]: + with self._lock: + self._check_open() + return self._backend.notifications() + + def close(self) -> None: + """Close in dependency order; a failure retains resources for retry. + + 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 + + +class Registration: + def __init__(self, endpoint: Endpoint, handle: Any, owner: Any): + self._endpoint, self._handle, self._owner = endpoint, handle, owner + + 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 Catalog: + def __init__(self, endpoint: Endpoint, handle: Any, peer: str | None): + self._endpoint, self._handle, self.peer = endpoint, handle, peer + + 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) + + +class Work: + def __init__( + 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 state(self) -> str: + """Last observed state without an extra provider query.""" + return self._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 + + 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: + """Release completed work; active cancellation is outside this MVP.""" + ep = self._endpoint + with ep._lock: + if self._closed: + return + 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 new file mode 100644 index 0000000000000..c7c900569d491 --- /dev/null +++ b/torch/distributed/_transfer/_backend.py @@ -0,0 +1,35 @@ +from typing import Any, Protocol + + +class Backend(Protocol): + """Provider contract; methods are called under the Endpoint's host lock. + + 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. + """ + + 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/_reference.py b/torch/distributed/_transfer/_reference.py new file mode 100644 index 0000000000000..70cd4e6a46571 --- /dev/null +++ b/torch/distributed/_transfer/_reference.py @@ -0,0 +1,94 @@ +"""Synchronous, in-process CPU backend for examples/tests, NOT networking.""" + +import ctypes +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