diff --git a/docs/nixl_ucx_payload.md b/docs/nixl_ucx_payload.md index 48beb1f2..c5a29d26 100644 --- a/docs/nixl_ucx_payload.md +++ b/docs/nixl_ucx_payload.md @@ -76,6 +76,14 @@ or GID, add the corresponding variables to `ucx_env_vars`. If NIXL initialization or a transfer fails, TQ reports the error directly and does not fall back to ZMQ. If the transport is not restricted, or if `UCX_TLS` includes `tcp`, UCX may use TCP. +### Receive Buffer Cache + +Set `payload_transfer.receive_buffer_cache_mb` to limit idle registered buffers **per runtime**, +in MB (1024² bytes). The default is `256`; `0` disables caching. Larger payloads still transfer, +but their buffers are released when no longer referenced. Live data, in-flight transfers and +fault-retained memory are outside this limit. A smaller cache saves idle memory at the cost of +more registrations; account for all runtimes on a node when increasing it. + ### Common UCX Configuration | Variable | Purpose | Reference value | @@ -117,6 +125,19 @@ validation only confirm that the NIXL-UCX path is usable; to confirm RDMA, also ## Common Issues +### A NIXL Session Has Failed + +An uncertain WRITE retains its buffers until teardown and stops the affected peer session. +Recreate the runtime through the owning worker's error handling. Confirmed pre-WRITE failures +and local decode errors after completion do not fail the GET session. + +PUT cancellation waits for its running send. GET cancellation after COMMIT waits within the +remaining timeout: a matching completion or pre-WRITE error releases the receiver; an unknown +result quarantines it and stops that GET target. Cancellation still propagates to the caller. + +With metrics enabled, `tq_storage_payload_*` reports cache, leased and quarantined bytes, +queue wait and sender resources. + ### RDMA Devices Are Ready, but NIXL-UCX Fails to Start If `ibv_devinfo` shows RDMA devices and active ports but NIXL initialization fails, the log typically contains: diff --git a/tests/test_metrics.py b/tests/test_metrics.py index eea13564..d5359bc3 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -457,6 +457,32 @@ def mock_query(su_info, su_id): assert exporter.storage_capacity.labels(storage_unit_id="SU_002")._value.get() == 500 assert call_count == 2 + def test_payload_metrics_collected_and_removed_when_backend_no_longer_reports_them(self): + exporter = TQMetricsExporter() + exporter._storage_unit_infos = {"storage": MagicMock()} + snapshot = { + "storage_unit_id": "storage", + "payload_transfer": {"idle_receive_bytes": 128, "queue_wait_seconds": 2.5}, + } + exporter._query_storage_unit = MagicMock(return_value=snapshot) + + exporter.collect_storage_metrics() + assert ( + exporter.registry.get_sample_value("tq_storage_payload_idle_receive_bytes", {"storage_unit_id": "storage"}) + == 128 + ) + assert ( + exporter.registry.get_sample_value("tq_storage_payload_queue_wait_seconds", {"storage_unit_id": "storage"}) + == 2.5 + ) + + snapshot.pop("payload_transfer") + exporter.collect_storage_metrics() + assert ( + exporter.registry.get_sample_value("tq_storage_payload_idle_receive_bytes", {"storage_unit_id": "storage"}) + is None + ) + def test_storage_metrics_skips_capacity_when_none(self): """When capacity is None (unlimited storage), capacity/utilization gauges must not be populated, but other metrics still are. @@ -567,11 +593,25 @@ def worker(self): @staticmethod def dispatch(worker, operation, body): + from queue import SimpleQueue + + from transfer_queue.storage.payload_transfer import DeferredResponse + from transfer_queue.storage.simple_storage import _drain_deferred_responses + unit, monitor = worker request = ZMQMessage.create(request_type=operation, sender_id="manager", body=body) socket = MagicMock() socket.recv_multipart.return_value = [b"manager"] + request.serialize() - unit._process_one_worker_request(socket, monitor) + measurements = {} + identity, response = unit._process_one_worker_request(socket, monitor, measurements) + if isinstance(response, DeferredResponse): + assert response.future.done() + completions = SimpleQueue() + completions.put((bytes(identity), response.future)) + _drain_deferred_responses(completions, socket, unit.storage_unit_id, measurements) + else: + socket.send_multipart([identity] + response.serialize(), copy=False) + assert measurements == {} sent = socket.send_multipart.call_args.args[0] assert sent[0] == b"manager" return ZMQMessage.deserialize(sent[1:]) @@ -609,7 +649,7 @@ def test_zmq_metrics_and_checkpoints_keep_their_arrival_labels(self, worker, tmp @pytest.mark.parametrize("finish", ["COMMIT", "CANCEL"]) def test_nixl_counts_prepare_once_and_skips_continuations(self, worker, operation, finish): from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer, PayloadDescriptor - from transfer_queue.utils.serial_utils import calc_packed_size, encode, pack_into + from transfer_queue.utils.serial_utils import encode unit, monitor = worker unit.storage_data.put_data({"value": [42]}, [0]) @@ -619,12 +659,9 @@ def test_nixl_counts_prepare_once_and_skips_continuations(self, worker, operatio transfer._runtime = MagicMock() unit.payload_transfer = transfer frames = encode({"value": [42]}) - payload = bytearray(calc_packed_size(frames)) - pack_into(payload, frames) - descriptor = PayloadDescriptor("payload", len(payload), tuple(memoryview(frame).nbytes for frame in frames)) - received = Future() - received.set_result(memoryview(payload)) - transfer._runtime.receive.return_value = received + frame_sizes = tuple(memoryview(frame).nbytes for frame in frames) + descriptor = PayloadDescriptor("payload", sum(frame_sizes), frame_sizes) + transfer._runtime.receive.return_value = tuple(memoryview(frame) for frame in frames) sent = Future() sent.set_result(None) transfer._runtime.send.return_value = sent diff --git a/tests/test_nixl_payload_lifecycle.py b/tests/test_nixl_payload_lifecycle.py new file mode 100644 index 00000000..b3b15629 --- /dev/null +++ b/tests/test_nixl_payload_lifecycle.py @@ -0,0 +1,1538 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team + +"""Focused tests for the frame-native NIXL payload lifecycle.""" + +from __future__ import annotations + +import asyncio +import gc +import socket +import sys +import threading +import time +import types +from concurrent.futures import Future +from contextlib import ExitStack +from queue import SimpleQueue +from typing import Any, Callable +from unittest.mock import AsyncMock, MagicMock + +import numpy as np +import pytest +import torch +import zmq + +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager +from transfer_queue.storage.payload_transfer import DeferredResponse, PayloadTransferError, create_payload_transfer +from transfer_queue.storage.payload_transfer.nixl import ( + NixlPayloadTransfer, + PayloadDescriptor, + ReceiveToken, + TransferEndpoint, + _PendingGet, +) +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import ( + NixlError, + NixlRuntime, + NixlWriteNotStarted, + _frame_regions, + _FrameSource, + _RegisteredReceiveBuffer, +) +from transfer_queue.storage.simple_storage import _drain_deferred_responses, _queue_deferred_response +from transfer_queue.utils.serial_utils import decode, encode +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType + + +class _FakeAgent: + def __init__(self) -> None: + self.backends = {"UCX": object()} + self.register_calls: list[Any] = [] + self.deregister_calls: list[Any] = [] + self.events: list[str] = [] + self.add_remote_calls: list[bytes] = [] + self.remove_remote_calls: list[str] = [] + self.transfer_calls: list[Any] = [] + self.metadata_names: dict[bytes, str] = {} + self.metadata_calls: list[Any] = [] + self.status_by_remote: dict[str, str] = {} + self.release_error = False + self.deregister_error = False + self._registration_id = 0 + self._handle_id = 0 + + def get_partial_agent_metadata(self, descs: Any, *, inc_conn_info: bool, backends: list[str]) -> bytes: + self.metadata_calls.append((descs, inc_conn_info, backends)) + return f"metadata-{descs[1]}".encode() + + def register_memory(self, regions: Any, **kwargs: Any) -> tuple[str, int]: + self.events.append("register") + self.register_calls.append(regions) + self._registration_id += 1 + return ("registration", self._registration_id) + + def deregister_memory(self, registration: Any, **kwargs: Any) -> None: + self.events.append("deregister") + if self.deregister_error: + raise RuntimeError("deregister failed") + self.deregister_calls.append(registration) + + def get_xfer_descs(self, regions: Any, **kwargs: Any) -> Any: + return list(regions) + + def get_serialized_descs(self, descs: Any) -> bytes: + return repr(descs).encode() + + def deserialize_descs(self, descs: bytes) -> bytes: + return descs + + def add_remote_agent(self, metadata: bytes) -> str: + self.add_remote_calls.append(metadata) + return self.metadata_names[metadata] + + def remove_remote_agent(self, remote_name: str) -> None: + self.remove_remote_calls.append(remote_name) + + def initialize_xfer(self, operation: str, local: Any, remote: Any, remote_name: str, **kwargs: Any) -> Any: + self._handle_id += 1 + return (self._handle_id, remote_name, local, remote) + + def transfer(self, handle: Any) -> str: + self.transfer_calls.append(handle) + return self.status_by_remote.get(handle[1], "DONE") + + def check_xfer_state(self, handle: Any) -> str: + return self.status_by_remote.get(handle[1], "DONE") + + def release_xfer_handle(self, handle: Any) -> None: + if self.release_error: + raise RuntimeError("release failed") + + +@pytest.fixture +def runtime(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest): + agents: list[_FakeAgent] = [] + module = types.ModuleType("nixl") + + def make_agent(*args: Any, **kwargs: Any) -> _FakeAgent: + agent = _FakeAgent() + agents.append(agent) + return agent + + module.nixl_agent = make_agent + module.nixl_agent_config = lambda **kwargs: kwargs + monkeypatch.setitem(sys.modules, "nixl", module) + instance = NixlRuntime() + # Keep lifecycle allocations tiny; None uses the public default. + cache_bytes = getattr(request, "param", None) + if cache_bytes is not None: + instance._receive_buffer_cache_bytes = cache_bytes + yield instance, agents[0] + instance.close() + + +def _descriptor(transfer_id: str, *sizes: int) -> PayloadDescriptor: + return PayloadDescriptor(transfer_id, sum(sizes), tuple(sizes)) + + +def _token(remote_name: str, metadata: bytes, descriptor: PayloadDescriptor) -> dict[str, Any]: + return { + "agent_name": remote_name, + "agent_metadata": metadata, + "frame_remote_descs": b"remote-descs", + "payload_bytes": descriptor.payload_bytes, + } + + +@pytest.mark.parametrize("configured, expected_mb", [(None, 256), (0, 0), (64, 64)]) +def test_receive_cache_config_uses_megabytes(runtime, configured, expected_mb): + config = {"backend": "nixl-ucx"} + if configured is not None: + config["receive_buffer_cache_mb"] = configured + transfer = create_payload_transfer(config) + try: + assert transfer.diagnostics["receive_buffer_cache_bytes"] == expected_mb * 1024**2 + finally: + transfer.close() + + +def test_frame_native_layout_preserves_empty_frames_and_large_offsets(): + sizes = (3 * 1024**3, 0, 3 * 1024**3, 2 * 1024**3) + descriptor = _descriptor("large", *sizes) + descriptor.validate() + + address = 0x1000 + regions = _frame_regions(address, descriptor.frame_sizes) + + assert descriptor.payload_bytes == 8 * 1024**3 + assert regions == [ + (address, 3 * 1024**3, 0), + (address + 3 * 1024**3, 3 * 1024**3, 0), + (address + 6 * 1024**3, 2 * 1024**3, 0), + ] + + +def test_all_empty_frames_skip_registration_and_write(runtime): + instance, agent = runtime + descriptor = _descriptor("empty", 0, 0) + + token = instance.prepare_receive(descriptor) + frames = instance.receive(descriptor) + sent = instance.send({"agent_name": "peer"}, token, descriptor, (b"", b"")) + + assert [bytes(frame) for frame in frames] == [b"", b""] + assert sent.result() is None + assert agent.register_calls == [] + assert agent.transfer_calls == [] + assert agent.metadata_calls == [] + assert token["agent_metadata"] == b"" + + +def test_receive_metadata_excludes_other_live_registrations(runtime): + instance, agent = runtime + leased = _descriptor("leased", 8) + instance.prepare_receive(leased) + held_frames = instance.receive(leased) + instance.prepare_receive(_descriptor("pending", 16)) + current = _descriptor("current", 32) + + token = instance.prepare_receive(current) + registration = instance._receives[current.transfer_id].registration + + assert len(agent.register_calls) == 3 + assert agent.metadata_calls[-1] == (registration, True, ["UCX"]) + assert token["agent_metadata"] == f"metadata-{registration[1]}".encode() + assert held_frames[0].nbytes == 8 + + +@pytest.mark.parametrize("runtime", [0, None], indirect=True) +def test_receive_frames_decode_without_copy_and_lease_until_last_tensor(runtime): + instance, agent = runtime + source = {"value": torch.arange(8, dtype=torch.int32)} + encoded = tuple(encode(source)) + descriptor = _descriptor("receive", *(memoryview(frame).nbytes for frame in encoded)) + + instance.prepare_receive(descriptor) + scratch = instance._receives[descriptor.transfer_id] + assert scratch is not None + offset = 0 + for frame in encoded: + view = memoryview(frame) + scratch.buffer[offset : offset + view.nbytes] = view + offset += view.nbytes + + frames = instance.receive(descriptor) + decoded = decode(list(frames)) + detached = decoded["value"].detach() + del frames, decoded + gc.collect() + + assert instance._idle_receive_buffers == [] + assert agent.deregister_calls == [] + assert torch.equal(detached, source["value"]) + + del detached + gc.collect() + if instance.diagnostics["receive_buffer_cache_bytes"] == 0: + assert instance._idle_receive_buffers == [] + assert agent.deregister_calls == [scratch.registration] + assert instance.diagnostics["registered_bytes"] == 0 + else: + assert instance._idle_receive_buffers == [scratch] + assert instance.diagnostics["idle_receive_bytes"] == scratch.capacity + + +@pytest.mark.parametrize("runtime", [64], indirect=True) +@pytest.mark.parametrize("capacities, required, reused_index", [((16,), 8, 0), ((16, 24), 12, 0), ((64,), 8, None)]) +def test_receive_pool_reuses_best_fit_but_rejects_oversized_buffers(runtime, capacities, required, reused_index): + instance, agent = runtime + for index, capacity in enumerate(capacities): + instance.prepare_receive(_descriptor(str(index), capacity)) + buffers = list(instance._receives.values()) + for index in range(len(capacities)): + instance.cancel_receive(str(index)) + + request = _descriptor("request", required) + instance.prepare_receive(request) + + if reused_index is None: + assert instance._receives[request.transfer_id].capacity == required + assert agent.deregister_calls == [buffers[0].registration] + else: + assert instance._receives[request.transfer_id] is buffers[reused_index] + assert agent.deregister_calls == [] + registrations = len(capacities) + (reused_index is None) + assert len(agent.register_calls) == registrations + assert instance.diagnostics["receive_buffer_acquisitions"] == len(capacities) + 1 + assert instance.diagnostics["receive_buffer_registrations"] == registrations + assert instance.diagnostics["receive_buffer_reuses"] == (reused_index is not None) + + +def test_receive_pool_tracks_pending_and_leased_working_set(runtime): + instance, _ = runtime + leased = _descriptor("leased", 16) + pending = _descriptor("pending", 24) + + instance.prepare_receive(leased) + held_frames = instance.receive(leased) + instance.prepare_receive(pending) + + diagnostics = instance.diagnostics + assert diagnostics["receive_working_set_hwm_count"] == 2 + assert diagnostics["receive_working_set_hwm_bytes"] == 40 + + instance.cancel_receive(pending.transfer_id) + del held_frames + gc.collect() + assert instance.diagnostics["idle_receive_bytes"] == 40 + assert instance.diagnostics["pending_receive_bytes"] == 0 + assert instance.diagnostics["leased_receive_bytes"] == 0 + + +@pytest.mark.parametrize("runtime", [32], indirect=True) +def test_receive_pool_evicts_oldest_idle_before_registering_on_miss(runtime): + instance, agent = runtime + leased = _descriptor("leased", 16) + instance.prepare_receive(leased) + leased_buffer = instance._receives[leased.transfer_id] + held_frames = instance.receive(leased) + + undersized = _descriptor("undersized", 8) + instance.prepare_receive(undersized) + undersized_buffer = instance._receives[undersized.transfer_id] + + larger_idle = _descriptor("larger_idle", 24) + instance.prepare_receive(larger_idle) + larger_idle_buffer = instance._receives[larger_idle.transfer_id] + instance.cancel_receive(undersized.transfer_id) + instance.cancel_receive(larger_idle.transfer_id) + + request = _descriptor("request", 32) + instance.prepare_receive(request) + + assert instance._leased_receive_buffers[id(leased_buffer)] is leased_buffer + assert instance._receives[request.transfer_id].capacity == 32 + assert agent.deregister_calls == [ + undersized_buffer.registration, + larger_idle_buffer.registration, + ] + assert agent.events[-3:] == ["deregister", "deregister", "register"] + assert instance._idle_receive_buffers == [] + assert instance.diagnostics["registered_bytes"] == 16 + 32 + assert instance.diagnostics["receive_working_set_hwm_count"] == 3 + assert instance.diagnostics["receive_working_set_hwm_bytes"] == 48 + assert instance.diagnostics["receive_buffer_acquisitions"] == 4 + assert instance.diagnostics["receive_buffer_registrations"] == 4 + assert instance.diagnostics["receive_buffer_reuses"] == 0 + assert instance.diagnostics["receive_buffer_evictions"] == 2 + + instance.cancel_receive(request.transfer_id) + del held_frames + + +@pytest.mark.parametrize("runtime", [16], indirect=True) +def test_receive_pool_releases_excess_on_return_without_another_prepare(runtime): + instance, agent = runtime + first = _descriptor("first", 16) + second = _descriptor("second", 16) + instance.prepare_receive(first) + instance.prepare_receive(second) + second_buffer = instance._receives[second.transfer_id] + + instance.cancel_receive(first.transfer_id) + instance.cancel_receive(second.transfer_id) + + assert instance.diagnostics["idle_receive_bytes"] == 16 + assert instance.diagnostics["registered_bytes"] == 16 + assert instance.diagnostics["pending_receive_buffers"] == 0 + assert agent.deregister_calls == [second_buffer.registration] + + +@pytest.mark.parametrize("runtime", [16], indirect=True) +def test_receive_larger_than_cache_preserves_live_data_and_warm_idle_buffer(runtime): + instance, agent = runtime + warm = _descriptor("warm", 8) + instance.prepare_receive(warm) + warm_buffer = instance._receives[warm.transfer_id] + instance.cancel_receive(warm.transfer_id) + + large = _descriptor("large", 32) + instance.prepare_receive(large) + large_buffer = instance._receives[large.transfer_id] + large_buffer.buffer[:] = b"x" * 32 + frames = instance.receive(large) + + assert bytes(frames[0]) == b"x" * 32 + assert instance._idle_receive_buffers == [warm_buffer] + assert agent.deregister_calls == [] + + del frames + gc.collect() + + assert agent.deregister_calls == [large_buffer.registration] + assert instance._idle_receive_buffers == [warm_buffer] + assert instance.diagnostics["idle_receive_bytes"] == 8 + assert instance.diagnostics["registered_bytes"] == 8 + + +@pytest.mark.parametrize("runtime", [0], indirect=True) +def test_receive_prepare_failure_releases_uncached_buffer(runtime, monkeypatch): + instance, agent = runtime + + def fail_metadata(*args, **kwargs): + raise RuntimeError("metadata failed") + + monkeypatch.setattr(agent, "get_partial_agent_metadata", fail_metadata) + with pytest.raises(NixlError, match="metadata failed"): + instance.prepare_receive(_descriptor("failed", 16)) + + assert instance._receives == {} + assert instance._idle_receive_buffers == [] + assert len(agent.deregister_calls) == 1 + assert instance.diagnostics["registered_bytes"] == 0 + + +@pytest.mark.parametrize("runtime", [0], indirect=True) +@pytest.mark.parametrize("failing_call", ["np.frombuffer", "_frame_views"]) +def test_receive_exposure_failure_keeps_pending_buffer_owned(runtime, monkeypatch, failing_call): + instance, agent = runtime + descriptor = _descriptor("pending", 16) + instance.prepare_receive(descriptor) + scratch = instance._receives[descriptor.transfer_id] + + def fail_owner(*args, **kwargs): + raise RuntimeError("view creation failed") + + monkeypatch.setattr(f"transfer_queue.storage.payload_transfer.nixl_ucx_runtime.{failing_call}", fail_owner) + with pytest.raises(NixlError, match="view creation failed"): + instance.receive(descriptor) + + assert instance._receives[descriptor.transfer_id] is scratch + assert instance._leased_receive_buffers == {} + assert agent.deregister_calls == [] + instance.cancel_receive(descriptor.transfer_id) + assert agent.deregister_calls == [scratch.registration] + assert instance.diagnostics["registered_bytes"] == 0 + + +@pytest.mark.parametrize("runtime", [0, 16], indirect=True) +def test_receive_cleanup_failure_retains_memory_and_stops_new_allocations(runtime): + instance, agent = runtime + descriptor = _descriptor("failed", 16) + instance.prepare_receive(descriptor) + scratch = instance._receives[descriptor.transfer_id] + agent.deregister_error = True + + instance.cancel_receive(descriptor.transfer_id) + + with pytest.raises(NixlError, match="receive memory cleanup failed"): + instance.prepare_receive(_descriptor("next", 4)) + assert instance._quarantined_receive_buffers == [scratch] + assert instance.diagnostics["quarantined_bytes"] == 16 + assert instance.diagnostics["registered_bytes"] == 16 + assert instance.diagnostics["idle_receive_bytes"] == 0 + assert len(agent.register_calls) == 1 + + +@pytest.mark.parametrize("runtime", [0], indirect=True) +def test_sources_reuse_leased_mr_and_external_registration_is_transfer_scoped(runtime): + instance, agent = runtime + descriptor = _descriptor("lease", 8) + instance.prepare_receive(descriptor) + leased_frame = instance.receive(descriptor)[0] + + reused = instance._acquire_frame_source(leased_frame) + writable = bytearray(b"writable") + direct = instance._acquire_frame_source(writable) + readonly = instance._acquire_frame_source(b"readonly") + non_contiguous = instance._acquire_frame_source(memoryview(np.arange(8, dtype=np.uint8))[::2]) + direct_again = instance._acquire_frame_source(writable) + + assert reused.registration is None + assert isinstance(direct.owner, memoryview) + assert isinstance(readonly.owner, bytearray) + assert isinstance(non_contiguous.owner, bytearray) + assert isinstance(direct_again.owner, bytearray) + assert direct_again.address != direct.address + assert direct.registration != direct_again.registration + assert len(agent.register_calls) == 5 # one receive MR plus four transfer-local registrations + + +def test_receive_metadata_reused_then_replaced_in_peer_executor(runtime): + instance, agent = runtime + descriptor = _descriptor("send", 1) + agent.metadata_names.update({b"m1": "peer", b"m2": "peer"}) + + instance.send({}, _token("peer", b"m1", descriptor), descriptor, (bytearray(b"a"),)).result() + instance.send({}, _token("peer", b"m1", descriptor), descriptor, (bytearray(b"b"),)).result() + instance.send({}, _token("peer", b"m2", descriptor), descriptor, (bytearray(b"c"),)).result() + + assert agent.add_remote_calls == [b"m1", b"m2"] + assert agent.remove_remote_calls == ["peer"] + + +def test_same_peer_serializes_while_different_peer_runs_concurrently(runtime, monkeypatch: pytest.MonkeyPatch): + instance, _ = runtime + descriptor = _descriptor("concurrency", 1) + first_started = threading.Event() + other_started = threading.Event() + release = threading.Event() + starts: list[str] = [] + + def blocking_send(remote_name: str, *args: Any) -> None: + starts.append(remote_name) + (first_started if remote_name == "peer-a" else other_started).set() + release.wait(timeout=2) + + monkeypatch.setattr(instance, "_send_scatter", blocking_send) + first = instance.send({}, _token("peer-a", b"a", descriptor), descriptor, (bytearray(b"a"),)) + assert first_started.wait(timeout=1) + queued = instance.send({}, _token("peer-a", b"a", descriptor), descriptor, (bytearray(b"b"),)) + other = instance.send({}, _token("peer-b", b"b", descriptor), descriptor, (bytearray(b"c"),)) + + assert other_started.wait(timeout=1) + assert starts.count("peer-a") == 1 + release.set() + first.result(timeout=1) + queued.result(timeout=1) + other.result(timeout=1) + assert starts.count("peer-a") == 2 + + +def test_completed_write_cleanup_failure_keeps_success_and_fails_peer(runtime): + instance, agent = runtime + descriptor = _descriptor("cleanup", 1) + agent.metadata_names[b"metadata"] = "peer" + agent.release_error = True + + completed = instance.send({}, _token("peer", b"metadata", descriptor), descriptor, (bytearray(b"x"),)) + + assert completed.result(timeout=1) is None + assert instance._retained_resources + with pytest.raises(NixlError, match="has failed"): + instance.send({}, _token("peer", b"metadata", descriptor), descriptor, (bytearray(b"y"),)) + + +def test_sender_reuses_idle_executor_without_rebinding_outstanding_peer(runtime, monkeypatch): + instance, _ = runtime + descriptor = _descriptor("reuse-executor", 1) + started = threading.Event() + release = threading.Event() + + def send(remote_name, *args): + if remote_name == "peer-a": + started.set() + release.wait(timeout=2) + + monkeypatch.setattr(instance, "_send_scatter", send) + first = instance.send({}, _token("peer-a", b"a", descriptor), descriptor, (b"x",)) + try: + assert started.wait(timeout=1) + first_sender = instance._peer_senders["peer-a"] + queued = instance.send({}, _token("peer-a", b"a", descriptor), descriptor, (b"x",)) + assert queued.cancel() + assert first_sender.outstanding == 1 + other = instance.send({}, _token("peer-b", b"b", descriptor), descriptor, (b"x",)) + assert instance._peer_senders["peer-b"] is not first_sender + other.result(timeout=1) + finally: + release.set() + first.result(timeout=1) + # A barrier also waits for the executor's Future completion callbacks. + first_sender.executor.submit(lambda: None).result(timeout=1) + + instance.send({}, _token("next-epoch", b"c", descriptor), descriptor, (b"x",)).result(timeout=1) + assert instance._peer_senders["next-epoch"] is first_sender + assert "peer-a" not in instance._peer_senders + assert instance.diagnostics["sender_executors"] == 2 + + +@pytest.mark.parametrize("outcome", ["success", "failure", "cancelled", "submit_error"]) +def test_sender_retires_immediate_completion_and_submit_failure(runtime, monkeypatch, outcome): + instance, agent = runtime + descriptor = _descriptor("retire", 1) + agent.metadata_names[b"metadata"] = "peer" + token = _token("peer", b"metadata", descriptor) + instance.send({}, token, descriptor, (b"x",)).result(timeout=1) + sender = instance._peer_senders["peer"] + sender.executor.submit(lambda: None).result(timeout=1) + completed = Future() + if outcome == "success": + completed.set_result(None) + elif outcome == "failure": + completed.set_exception(NixlError("write failed")) + else: + completed.cancel() + + def submit(*args): + if outcome == "submit_error": + raise RuntimeError("submit failed") + return completed + + monkeypatch.setattr(sender.executor, "submit", submit) + if outcome == "submit_error": + with pytest.raises(RuntimeError, match="submit failed"): + instance.send({}, token, descriptor, (b"x",)) + else: + assert instance.send({}, token, descriptor, (b"x",)) is completed + assert instance.diagnostics["outstanding_sends"] == 0 + + +def test_send_deadline_covers_queue_and_source_preparation(runtime, monkeypatch): + instance, agent = runtime + descriptor = _descriptor("deadline", 1) + agent.metadata_names[b"metadata"] = "peer" + token = _token("peer", b"metadata", descriptor) + clock = [0.0] + monkeypatch.setattr( + "transfer_queue.storage.payload_transfer.nixl_ucx_runtime.time", + types.SimpleNamespace(monotonic=lambda: clock[0], sleep=time.sleep), + ) + started = threading.Event() + release = threading.Event() + acquire = instance._acquire_frame_source + + def prepare(frame): + started.set() + release.wait(timeout=2) + return acquire(frame) + + monkeypatch.setattr(instance, "_acquire_frame_source", prepare) + first = instance.send({}, token, descriptor, (b"x",)) + try: + assert started.wait(timeout=1) + queued = instance.send({}, token, descriptor, (b"y",)) + clock[0] = instance._timeout_seconds + 1 + finally: + release.set() + for future in (first, queued): + with pytest.raises(NixlWriteNotStarted, match="including queue wait"): + future.result(timeout=1) + assert len(agent.register_calls) == len(agent.deregister_calls) == 1 + assert agent.transfer_calls == [] + assert instance._retained_resources == [] + assert instance._failed_peers == set() + assert instance.diagnostics["queue_wait_seconds"] >= instance._timeout_seconds + assert instance.send({}, token, descriptor, (b"z",)).result(timeout=1) is None + + +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_expired_prepared_send_releases_unposted_handle_without_write(runtime, monkeypatch, cleanup_fails): + instance, agent = runtime + descriptor = _descriptor("pre-write-timeout", 1) + agent.metadata_names[b"metadata"] = "peer" + agent.release_error = cleanup_fails + initialize = agent.initialize_xfer + check_deadline = instance._check_send_deadline + + def expire(deadline): + raise NixlError("send deadline expired") + + def initialize_then_expire(*args, **kwargs): + handle = initialize(*args, **kwargs) + monkeypatch.setattr(instance, "_check_send_deadline", expire) + return handle + + monkeypatch.setattr(agent, "initialize_xfer", initialize_then_expire) + token = _token("peer", b"metadata", descriptor) + with pytest.raises(NixlWriteNotStarted, match="deadline expired"): + instance.send({}, token, descriptor, (b"x",)).result(timeout=1) + assert agent.transfer_calls == [] + assert ("peer" in instance._failed_peers) == cleanup_fails + assert bool(instance._retained_resources) == cleanup_fails + if not cleanup_fails: + monkeypatch.setattr(instance, "_check_send_deadline", check_deadline) + monkeypatch.setattr(agent, "initialize_xfer", initialize) + assert instance.send({}, token, descriptor, (b"x",)).result(timeout=1) is None + + +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_local_preparation_failure_only_fails_peer_if_cleanup_fails(runtime, monkeypatch, cleanup_fails): + instance, agent = runtime + descriptor = _descriptor("prepare-failure", 1, 1) + agent.metadata_names[b"metadata"] = "peer" + agent.deregister_error = cleanup_fails + acquire = instance._acquire_frame_source + calls = 0 + + def prepare(frame): + nonlocal calls + calls += 1 + if calls == 2: + raise MemoryError("source allocation failed") + return acquire(frame) + + monkeypatch.setattr(instance, "_acquire_frame_source", prepare) + token = _token("peer", b"metadata", descriptor) + with pytest.raises(NixlWriteNotStarted, match="source allocation failed"): + instance.send({}, token, descriptor, (b"x", b"y")).result(timeout=1) + assert agent.transfer_calls == [] + assert ("peer" in instance._failed_peers) == cleanup_fails + assert bool(instance._retained_resources) == cleanup_fails + if not cleanup_fails: + assert instance.send({}, token, descriptor, (b"x", b"y")).result(timeout=1) is None + + +def test_metadata_failure_still_fails_peer_before_handle_creation(runtime): + instance, _ = runtime + descriptor = _descriptor("bad-metadata", 1) + token = _token("peer", b"unknown-metadata", descriptor) + with pytest.raises(NixlWriteNotStarted): + instance.send({}, token, descriptor, (b"x",)).result(timeout=1) + with pytest.raises(NixlError, match="has failed"): + instance.send({}, token, descriptor, (b"x",)) + + +def test_failed_write_retains_resources_and_other_peer_still_works(runtime): + instance, agent = runtime + descriptor = _descriptor("failure", 1) + agent.metadata_names.update({b"bad": "bad-peer", b"good": "good-peer"}) + agent.status_by_remote["bad-peer"] = "ERR" + + with pytest.raises(NixlError, match="status"): + instance.send({}, _token("bad-peer", b"bad", descriptor), descriptor, (bytearray(b"x"),)).result() + + assert instance._retained_resources + assert instance.send({}, _token("good-peer", b"good", descriptor), descriptor, (bytearray(b"y"),)).result() is None + + +def test_failed_future_does_not_retain_handle_or_sources(runtime): + instance, agent = runtime + descriptor = _descriptor("failure-traceback", 1) + agent.metadata_names[b"metadata"] = "peer" + agent.status_by_remote["peer"] = "ERR" + + future = instance.send({}, _token("peer", b"metadata", descriptor), descriptor, (bytearray(b"x"),)) + error = future.exception(timeout=1) + + assert isinstance(error, NixlError) + assert not isinstance(error, NixlWriteNotStarted) + traceback = error.__traceback__ + while traceback is not None and traceback.tb_frame.f_code.co_name != "_send_scatter": + traceback = traceback.tb_next + assert traceback is not None + assert traceback.tb_frame.f_locals["handle"] is None + assert traceback.tb_frame.f_locals["sources"] == [] + + +def test_close_interrupts_active_polling(runtime): + instance, agent = runtime + descriptor = _descriptor("closing", 1) + agent.metadata_names[b"metadata"] = "peer" + agent.status_by_remote["peer"] = "PROC" + future = instance.send({}, _token("peer", b"metadata", descriptor), descriptor, (bytearray(b"x"),)) + + deadline = time.monotonic() + 1 + while not agent.transfer_calls and time.monotonic() < deadline: + time.sleep(0.001) + started = time.monotonic() + instance.close() + + assert time.monotonic() - started < 1 + with pytest.raises(NixlError, match="closing"): + future.result() + + +def test_close_tears_down_agent_before_retained_owners(runtime): + instance, _ = runtime + events: list[str] = [] + + class TrackingAgent: + def __del__(self) -> None: + events.append("agent") + + class AgentHandle: + def __init__(self, agent: TrackingAgent, release_lease: Callable[[], None]) -> None: + self.agent = agent + self.release_lease = release_lease + + def __del__(self) -> None: + events.append("handle") + self.release_lease() + + class TrackingBuffer(bytearray): + def __init__(self, name: str) -> None: + super().__init__(b"x") + self.name = name + + def __del__(self) -> None: + events.append(self.name) + + agent = TrackingAgent() + source = _FrameSource(TrackingBuffer("source"), None, 0, 1) + receiver = _RegisteredReceiveBuffer(TrackingBuffer("receiver"), object(), 0) + leased_receiver = _RegisteredReceiveBuffer(TrackingBuffer("leased_receiver"), object(), 1) + lease_key = id(leased_receiver) + handle = AgentHandle(agent, lambda: instance._release_lease(lease_key)) + instance._agent = agent + instance._retained_resources = [(handle, [source])] + instance._quarantined_receive_buffers = [receiver] + instance._leased_receive_buffers[lease_key] = leased_receiver + del agent, handle, source, receiver, leased_receiver + + instance.close() + gc.collect() + + assert events.index("handle") < events.index("agent") + assert events.index("agent") < events.index("source") + assert events.index("agent") < events.index("receiver") + assert events.index("agent") < events.index("leased_receiver") + + +def test_quarantine_never_returns_exposed_receiver_to_pool(runtime): + instance, _ = runtime + descriptor = _descriptor("unsafe", 32) + instance.prepare_receive(descriptor) + scratch = instance._receives["unsafe"] + + instance.quarantine_receive("unsafe") + + assert instance._idle_receive_buffers == [] + assert instance._quarantined_receive_buffers == [scratch] + assert instance.diagnostics["quarantined_bytes"] == 32 + + safe = _descriptor("safe", 8) + instance.prepare_receive(safe) + assert instance.diagnostics["receive_working_set_hwm_count"] == 1 + assert instance.diagnostics["receive_working_set_hwm_bytes"] == 32 + assert instance.diagnostics["registered_bytes"] == 40 + + +@pytest.mark.parametrize("outcome", ["done", "cancelled", "not_started", "failed"]) +def test_get_commit_defers_response_until_send_finishes(outcome): + transfer = object.__new__(NixlPayloadTransfer) + descriptor = _descriptor("get", 1) + transfer._pending_gets = {"get": _PendingGet(descriptor, "manager", (bytearray(b"x"),))} + send_future: Future[None] = Future() + transfer.send = lambda *args, **kwargs: send_future + request = ZMQMessage.create( + request_type=ZMQRequestType.GET_DATA_COMMIT, + sender_id="manager", + body={ + "transfer_id": "get", + "receiver_endpoint": {"transport": "nixl-ucx", "data": {}}, + "receive_token": {"data": {}}, + }, + ) + + response = transfer._handle_get_commit(request, "storage") + + assert isinstance(response, DeferredResponse) + assert not response.future.done() + if outcome == "done": + send_future.set_result(None) + elif outcome == "cancelled": + send_future.cancel() + else: + error = NixlWriteNotStarted("not posted") if outcome == "not_started" else NixlError("write failed") + send_future.set_exception(error) + + result = response.future.result() + assert result.body["transfer_id"] == "get" + assert result.request_type == ( + ZMQRequestType.GET_DATA_RESPONSE if outcome == "done" else ZMQRequestType.PUT_GET_ERROR + ) + if outcome != "done": + assert result.body["write_not_started"] == (outcome in ("cancelled", "not_started")) + + +class _ProtocolSocket: + def __init__(self, *, fail_commit: bool = False, fail_ready: bool = False) -> None: + self.fail_commit = fail_commit + self.fail_ready = fail_ready + self.last_request: ZMQMessage | None = None + + async def send_multipart(self, frames: Any, **kwargs: Any) -> None: + request = ZMQMessage.deserialize(frames) + if self.fail_commit and request.request_type in ( + ZMQRequestType.GET_DATA_COMMIT, + ZMQRequestType.PUT_DATA_COMMIT, + ): + raise RuntimeError("commit send failed") + self.last_request = request + + async def recv_multipart(self, **kwargs: Any) -> Any: + assert self.last_request is not None + request = self.last_request + if request.request_type == ZMQRequestType.PUT_DATA_PREPARE: + if self.fail_ready: + raise RuntimeError("ready receive failed") + response = ZMQMessage.create( + request_type=ZMQRequestType.PUT_DATA_READY, + sender_id="storage", + body={ + "descriptor": request.body["descriptor"], + "receive_token": {"data": {}}, + }, + ) + elif request.request_type == ZMQRequestType.GET_DATA_PREPARE: + descriptor = _descriptor(str(request.body["transfer_id"]), 1) + response = ZMQMessage.create( + request_type=ZMQRequestType.GET_DATA_READY, + sender_id="storage", + body={"descriptor": descriptor.to_dict()}, + ) + elif request.request_type in (ZMQRequestType.PUT_DATA_COMMIT, ZMQRequestType.GET_DATA_COMMIT): + response = ZMQMessage.create( + request_type=( + ZMQRequestType.PUT_DATA_RESPONSE + if request.request_type == ZMQRequestType.PUT_DATA_COMMIT + else ZMQRequestType.GET_DATA_RESPONSE + ), + sender_id=request.receiver_id, + body={"transfer_id": request.body["transfer_id"]}, + ) + else: + raise AssertionError(f"unexpected receive after {request.request_type}") + return response.serialize() + + +class _LoopbackSocket: + def __init__(self, storage: NixlPayloadTransfer): + self.storage = storage + + async def send_multipart(self, frames, **kwargs): + self.response = self.storage.handle_request( + ZMQMessage.deserialize(frames), + storage_id="storage", + load_data=lambda *args: {"value": [torch.arange(8)]}, + store_data=lambda *args: pytest.fail("failed PUT must not publish data"), + ) + + async def recv_multipart(self, **kwargs): + response = self.response + if isinstance(response, DeferredResponse): + response = await asyncio.wrap_future(response.future) + return response.serialize() + + +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("failure_stage", ["prepare", "write"]) +def test_prewrite_failure_releases_receiver_across_protocol_and_runtime(runtime, operation, failure_stage): + storage = NixlPayloadTransfer(receive_buffer_cache_mb=0) + client = NixlPayloadTransfer(peer_infos={"storage": storage.bootstrap_info()}, receive_buffer_cache_mb=0) + sender = client if operation == "put" else storage + sender._runtime._timeout_seconds = 0 + + class FailingSocket(_LoopbackSocket): + async def recv_multipart(self, **kwargs): + if failure_stage == "prepare" and self.response.request_type in ( + ZMQRequestType.PUT_DATA_READY, + ZMQRequestType.GET_DATA_READY, + ): + raise zmq.Again() + return await super().recv_multipart(**kwargs) + + async def exercise(): + for _ in range(2): + kwargs = {"data": {"value": [1]}, "data_parser": None} if operation == "put" else {"fields": ["value"]} + with pytest.raises(RuntimeError, match="timed out|recv timeout"): + await getattr(client, operation)( + control_socket=FailingSocket(storage), + sender_id="manager", + target_id="storage", + global_indexes=[1], + **kwargs, + ) + assert storage._pending_puts == storage._pending_gets == {} + for transfer in (storage, client): + assert transfer._runtime._receives == {} + assert transfer.diagnostics["registered_bytes"] == 0 + assert transfer.diagnostics["quarantined_bytes"] == 0 + assert client._failed_get_targets == set() + assert sender._runtime._agent.transfer_calls == [] + assert sender._runtime._failed_peers == set() + + try: + asyncio.run(exercise()) + finally: + client.close() + storage.close() + + +@pytest.mark.parametrize("failing_call", ["np.frombuffer", "_frame_views"]) +def test_put_commit_releases_receive_when_view_creation_fails(runtime, monkeypatch, failing_call): + storage = NixlPayloadTransfer(receive_buffer_cache_mb=0) + descriptor = _descriptor("failed-view", 16) + prepare = ZMQMessage.create( + request_type=ZMQRequestType.PUT_DATA_PREPARE, + sender_id="manager", + body={"descriptor": descriptor.to_dict(), "global_indexes": [1]}, + ) + commit = ZMQMessage.create( + request_type=ZMQRequestType.PUT_DATA_COMMIT, + sender_id="other-manager", + body={"transfer_id": descriptor.transfer_id}, + ) + + def fail_owner(*args, **kwargs): + raise MemoryError("view allocation failed") + + try: + assert storage._handle_put_prepare(prepare, "storage").request_type == ZMQRequestType.PUT_DATA_READY + store_data = lambda *args: pytest.fail("failed PUT must not publish data") + assert storage._handle_put_commit(commit, "storage", store_data).request_type == ZMQRequestType.PUT_GET_ERROR + assert descriptor.transfer_id in storage._pending_puts + assert descriptor.transfer_id in storage._runtime._receives + + monkeypatch.setattr(f"transfer_queue.storage.payload_transfer.nixl_ucx_runtime.{failing_call}", fail_owner) + commit = ZMQMessage.create( + request_type=ZMQRequestType.PUT_DATA_COMMIT, + sender_id="manager", + body={"transfer_id": descriptor.transfer_id}, + ) + response = storage._handle_put_commit(commit, "storage", store_data) + + assert response.request_type == ZMQRequestType.PUT_GET_ERROR + assert "view allocation failed" in response.body["message"] + assert storage._pending_puts == storage._runtime._receives == {} + assert storage.diagnostics["registered_bytes"] == 0 + finally: + storage.close() + + +@pytest.mark.parametrize("error_type", [NixlError, NixlWriteNotStarted]) +def test_put_only_cancels_failed_send_when_write_never_started(error_type): + transfer = object.__new__(NixlPayloadTransfer) + failed_send: Future[None] = Future() + failed_send.set_exception(error_type("write failed")) + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + transfer.send = lambda *args, **kwargs: failed_send + cancellations: list[Any] = [] + + async def cancel(*args: Any) -> None: + cancellations.append(args) + + transfer._cancel = cancel + with pytest.raises(NixlError, match="write failed"): + asyncio.run( + transfer.put( + control_socket=_ProtocolSocket(), + sender_id="manager", + target_id="storage", + global_indexes=[1], + data={"value": [1]}, + data_parser=None, + ) + ) + + assert bool(cancellations) == (error_type is NixlWriteNotStarted) + + +@pytest.mark.parametrize("failure_stage", ["ready", "endpoint", "submit"]) +def test_put_cancels_prepared_receiver_before_write(failure_stage): + transfer = object.__new__(NixlPayloadTransfer) + control_socket = _ProtocolSocket(fail_ready=failure_stage == "ready") + + def fail(*args): + raise NixlError(f"{failure_stage} failed") + + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + if failure_stage == "endpoint": + transfer._peer_endpoint = fail + elif failure_stage == "submit": + transfer.send = fail + cancellations: list[Any] = [] + + async def cancel(*args: Any) -> None: + cancellations.append(args) + + transfer._cancel = cancel + with pytest.raises(RuntimeError, match=failure_stage): + asyncio.run( + transfer.put( + control_socket=control_socket, + sender_id="manager", + target_id="storage", + global_indexes=[1], + data={"value": [1]}, + data_parser=None, + ) + ) + + assert cancellations[0][2] == ZMQRequestType.PUT_DATA_CANCEL + assert cancellations[0][-1] is control_socket + + +@pytest.mark.parametrize("send_state", ["done", "queued", "running_done", "running_unposted", "running_failed"]) +def test_put_cancellation_at_write_completion_only_cancels_safe_receives(send_state): + transfer = object.__new__(NixlPayloadTransfer) + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + send_future = Future() + if send_state == "done": + send_future.set_result(None) + elif send_state.startswith("running"): + send_future.set_running_or_notify_cancel() + cancellations = [] + control_socket = _ProtocolSocket() + + def send(*args): + loop = asyncio.get_running_loop() + loop.call_soon(asyncio.current_task().cancel) + if send_state == "running_done": + loop.call_later(0.01, send_future.set_result, None) + elif send_state == "running_unposted": + loop.call_later(0.01, send_future.set_exception, NixlWriteNotStarted("not posted")) + elif send_state == "running_failed": + loop.call_later(0.01, send_future.set_exception, NixlError("uncertain write")) + return send_future + + async def cancel(*args): + assert send_future.done() + cancellations.append(args) + + transfer.send = send + transfer._cancel = cancel + with pytest.raises(asyncio.CancelledError): + asyncio.run( + transfer.put( + control_socket=control_socket, + sender_id="manager", + target_id="storage", + global_indexes=[1], + data={"value": [1]}, + data_parser=None, + ) + ) + assert send_future.done() + assert bool(cancellations) == (send_state != "running_failed") + if cancellations: + assert cancellations[0][-1] is control_socket + assert control_socket.last_request.request_type == ZMQRequestType.PUT_DATA_PREPARE + + +def test_put_does_not_cancel_after_commit_attempt(): + transfer = object.__new__(NixlPayloadTransfer) + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + completed = Future() + completed.set_result(None) + transfer.send = lambda *args: completed + + async def cancel(*args): + pytest.fail("COMMIT may already have published the data") + + transfer._cancel = cancel + with pytest.raises(RuntimeError, match="commit send failed"): + asyncio.run( + transfer.put( + control_socket=_ProtocolSocket(fail_commit=True), + sender_id="manager", + target_id="storage", + global_indexes=[1], + data={"value": [1]}, + data_parser=None, + ) + ) + + +@pytest.mark.parametrize("operation", ["PUT", "GET"]) +def test_cancel_drains_late_ready_and_matches_ack_on_original_socket(operation): + transfer = object.__new__(NixlPayloadTransfer) + replies = [ + ZMQMessage.create(request_type=ZMQRequestType[f"{operation}_DATA_READY"], sender_id="storage", body={}), + ZMQMessage.create( + request_type=ZMQRequestType[f"{operation}_DATA_RESPONSE"], + sender_id="storage", + body={"transfer_id": "another-transfer"}, + ), + ZMQMessage.create( + request_type=ZMQRequestType[f"{operation}_DATA_RESPONSE"], + sender_id="another-storage", + body={"transfer_id": "cancelled-transfer"}, + ), + ZMQMessage.create( + request_type=ZMQRequestType[f"{operation}_DATA_RESPONSE"], + sender_id="storage", + body={"transfer_id": "cancelled-transfer"}, + ), + ] + + class CancelSocket(_ProtocolSocket): + async def recv_multipart(self, **kwargs): + return replies.pop(0).serialize() + + control_socket = CancelSocket() + confirmed = asyncio.run( + transfer._cancel( + "manager", "storage", ZMQRequestType[f"{operation}_DATA_CANCEL"], "cancelled-transfer", control_socket + ) + ) + assert confirmed is True + assert replies == [] + assert control_socket.last_request.request_type == ZMQRequestType[f"{operation}_DATA_CANCEL"] + + +@pytest.mark.parametrize("operation", ["PUT", "GET"]) +@pytest.mark.parametrize("failure_stage", ["send", "receive"]) +def test_cancel_returns_false_when_ack_is_missing(operation, failure_stage): + transfer = object.__new__(NixlPayloadTransfer) + control_socket = MagicMock() + control_socket.send_multipart = AsyncMock(side_effect=zmq.Again() if failure_stage == "send" else None) + control_socket.recv_multipart = AsyncMock(side_effect=zmq.Again()) + + confirmed = asyncio.run( + transfer._cancel( + "manager", "storage", ZMQRequestType[f"{operation}_DATA_CANCEL"], "cancelled-transfer", control_socket + ) + ) + + assert confirmed is False + + +def test_get_cancels_remote_source_when_local_receive_prepare_fails(): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = set() + + def failed_prepare(descriptor: PayloadDescriptor) -> ReceiveToken: + raise NixlError("registration failed") + + transfer.prepare_receive = failed_prepare + cancellations: list[Any] = [] + + async def cancel(*args: Any) -> None: + cancellations.append(args) + + transfer._cancel = cancel + with pytest.raises(NixlError, match="registration failed"): + asyncio.run( + transfer.get( + control_socket=_ProtocolSocket(), + sender_id="manager", + target_id="storage", + global_indexes=[1], + fields=["value"], + ) + ) + + assert cancellations[0][2] == ZMQRequestType.GET_DATA_CANCEL + + +def test_get_commit_attempt_failure_quarantines_without_cancel(): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = set() + transfer.prepare_receive = lambda descriptor: ReceiveToken({"agent_name": "manager"}) + quarantined: list[str] = [] + cancellations: list[Any] = [] + transfer.quarantine_receive = quarantined.append + transfer.cancel_receive = lambda transfer_id: pytest.fail("exposed receiver was cancelled") + + async def cancel(*args: Any) -> None: + cancellations.append(args) + + transfer._cancel = cancel + with pytest.raises(RuntimeError, match="commit send failed"): + asyncio.run( + transfer.get( + control_socket=_ProtocolSocket(fail_commit=True), + sender_id="manager", + target_id="storage", + global_indexes=[1], + fields=["value"], + ) + ) + + assert len(quarantined) == 1 + assert cancellations == [] + assert transfer._failed_get_targets == {"storage"} + + retry_socket = _ProtocolSocket() + with pytest.raises(NixlError, match="GET session.*has failed"): + asyncio.run( + transfer.get( + control_socket=retry_socket, + sender_id="manager", + target_id="storage", + global_indexes=[1], + fields=["value"], + ) + ) + assert retry_socket.last_request is None + + +@pytest.mark.parametrize("cancel_stage", ["send", "receive"]) +@pytest.mark.parametrize( + "outcome", ["done", "not_started", "write_failed", "wrong_sender", "wrong_transfer", "timeout"] +) +def test_cancelled_get_reclaims_only_after_matching_write_completion(monkeypatch, cancel_stage, outcome): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = {"unrelated-target"} + transfer.prepare_receive = lambda descriptor: ReceiveToken({"agent_name": "manager"}) + transfer.receive = lambda descriptor: pytest.fail("cancelled GET must not expose its receive buffer") + released, quarantined = [], [] + transfer.cancel_receive = released.append + transfer.quarantine_receive = quarantined.append + if outcome == "timeout": + monkeypatch.setattr("transfer_queue.storage.payload_transfer.nixl.TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 0) + + class PendingSocket(_ProtocolSocket): + def __init__(self): + super().__init__() + self.started = asyncio.Event() + self.completion = asyncio.get_running_loop().create_future() + self.completion_reads = 0 + + async def send_multipart(self, frames, **kwargs): + await super().send_multipart(frames, **kwargs) + if self.last_request.request_type == ZMQRequestType.GET_DATA_COMMIT and cancel_stage == "send": + self.started.set() + await asyncio.get_running_loop().create_future() + + def recv_multipart(self, **kwargs): + if self.last_request.request_type == ZMQRequestType.GET_DATA_COMMIT: + self.completion_reads += 1 + self.started.set() + return self.completion + return super().recv_multipart(**kwargs) + + async def exercise(): + control_socket = PendingSocket() + task = asyncio.create_task( + transfer.get( + control_socket=control_socket, + sender_id="manager", + target_id="storage", + global_indexes=[1], + fields=["value"], + ) + ) + await control_socket.started.wait() + transfer_id = control_socket.last_request.body["transfer_id"] + task.cancel() + await asyncio.sleep(0) + if outcome != "timeout": + assert not task.done() + assert not control_socket.completion.cancelled() + is_error = outcome in ("not_started", "write_failed") + response = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR if is_error else ZMQRequestType.GET_DATA_RESPONSE, + sender_id="another-storage" if outcome == "wrong_sender" else "storage", + body={ + "transfer_id": "another-transfer" if outcome == "wrong_transfer" else transfer_id, + "write_not_started": outcome == "not_started", + }, + ) + control_socket.completion.set_result(response.serialize()) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(task), 1) + assert task.cancelled() + assert control_socket.completion_reads == 1 + assert control_socket.completion.done() + assert control_socket.last_request.request_type == ZMQRequestType.GET_DATA_COMMIT + if outcome in ("done", "not_started"): + assert released == [transfer_id] + assert quarantined == [] + assert transfer._failed_get_targets == {"unrelated-target"} + else: + assert released == [] + assert quarantined == [transfer_id] + assert transfer._failed_get_targets == {"unrelated-target", "storage"} + + asyncio.run(exercise()) + + +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("failure_stage", ["prepare", "commit"]) +def test_control_socket_errors_keep_runtime_error_contract(operation, failure_stage): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = set() + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + transfer.prepare_receive = lambda descriptor: ReceiveToken({"agent_name": "manager"}) + transfer.quarantine_receive = lambda transfer_id: None + completed = Future() + completed.set_result(None) + transfer.send = lambda *args: completed + + async def cancel(*args): + pass + + transfer._cancel = cancel + socket_error = zmq.Again() + + class FailingSocket(_ProtocolSocket): + async def recv_multipart(self, **kwargs): + if failure_stage == "prepare" or self.last_request.request_type in ( + ZMQRequestType.PUT_DATA_COMMIT, + ZMQRequestType.GET_DATA_COMMIT, + ): + raise socket_error + return await super().recv_multipart(**kwargs) + + kwargs = {"data": {"value": [1]}, "data_parser": None} if operation == "put" else {"fields": ["value"]} + with pytest.raises(PayloadTransferError, match=f"{operation.upper()}.*storage") as raised: + asyncio.run( + getattr(transfer, operation)( + control_socket=FailingSocket(), + sender_id="manager", + target_id="storage", + global_indexes=[1], + **kwargs, + ) + ) + assert isinstance(raised.value, RuntimeError) + assert raised.value.__cause__ is socket_error + + +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("timeout_on", ["send", "receive"]) +@pytest.mark.parametrize( + "failure_stage, cancel_confirmed, expected_attempts", + [("prepare", True, 2), ("prepare", False, 1), ("commit", True, 1)], +) +def test_manager_retries_only_confirmed_nixl_prepare_timeouts( + operation, timeout_on, failure_stage, cancel_confirmed, expected_attempts +): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = set() + transfer._peer_endpoint = lambda target_id: TransferEndpoint("nixl-ucx", {}) + transfer.prepare_receive = MagicMock(return_value=ReceiveToken({"agent_name": "manager"})) + transfer.receive = MagicMock(return_value=tuple(encode({"value": [1]}))) + transfer.cancel_receive = MagicMock() + transfer.quarantine_receive = MagicMock() + transfer._cancel = AsyncMock(return_value=cancel_confirmed) + completed = Future() + completed.set_result(None) + transfer.send = MagicMock(return_value=completed) + timeout = zmq.Again() + failing_request_type = ZMQRequestType[f"{operation.upper()}_DATA_{failure_stage.upper()}"] + sockets = [] + + class FailingSocket(_ProtocolSocket): + async def send_multipart(self, frames, **kwargs): + request = ZMQMessage.deserialize(frames) + if len(sockets) == 1 and timeout_on == "send" and request.request_type == failing_request_type: + raise timeout + await super().send_multipart(frames, **kwargs) + + async def recv_multipart(self, **kwargs): + if len(sockets) == 1 and timeout_on == "receive" and self.last_request.request_type == failing_request_type: + raise timeout + return await super().recv_multipart(**kwargs) + + manager = object.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "manager" + manager.storage_unit_infos = {} + manager.payload_transfer = transfer + manager.close = lambda: None + + async def attempt(): + control_socket = FailingSocket() + sockets.append(control_socket) + # Bypass socket leasing only; keep the manager's timeout conversion and retry loop. + if operation == "put": + return await manager._put_to_single_storage_unit.__wrapped__( + manager, [1], {"value": [1]}, "storage", socket=control_socket + ) + return await manager._get_from_single_storage_unit.__wrapped__( + manager, [1], ["value"], "storage", socket=control_socket + ) + + async def exercise(): + if expected_attempts == 2: + result = await manager._request_with_retry(operation, "storage", "samples=1", attempt, max_attempts=3) + assert result == (None if operation == "put" else (["value"], {"value": [1]})) + else: + message = "commit outcome is unknown" if failure_stage == "commit" else "without confirmed cancellation" + with pytest.raises(RuntimeError, match=message) as raised: + await manager._request_with_retry(operation, "storage", "samples=1", attempt, max_attempts=3) + payload_error = raised.value.__cause__ if operation == "put" else raised.value + assert type(payload_error) is PayloadTransferError + assert payload_error.__cause__ is timeout + + asyncio.run(exercise()) + + assert len(sockets) == expected_attempts + if failure_stage == "prepare": + transfer._cancel.assert_awaited_once() + assert transfer._cancel.await_args.args[-1] is sockets[0] + else: + transfer._cancel.assert_not_awaited() + expected_sends = int(operation == "put" and (expected_attempts == 2 or failure_stage == "commit")) + assert transfer.send.call_count == expected_sends + transfer.cancel_receive.assert_not_called() + if operation == "get" and failure_stage == "commit": + transfer.quarantine_receive.assert_called_once() + transfer.receive.assert_not_called() + assert transfer._failed_get_targets == {"storage"} + else: + transfer.quarantine_receive.assert_not_called() + assert transfer._failed_get_targets == set() + + +@pytest.mark.parametrize("failure_stage", ["receive", "decode"]) +def test_get_local_failure_after_ack_does_not_quarantine_or_fail_target(monkeypatch, failure_stage): + transfer = object.__new__(NixlPayloadTransfer) + transfer._failed_get_targets = {"unrelated-target"} + transfer.prepare_receive = lambda descriptor: ReceiveToken({"agent_name": "manager"}) + cancellations = [] + transfer.cancel_receive = cancellations.append + transfer.quarantine_receive = lambda transfer_id: pytest.fail("WRITE was already confirmed complete") + + def receive(descriptor): + if failure_stage == "receive": + raise NixlError("local receive failed") + return (memoryview(b"x"),) + + transfer.receive = receive + if failure_stage == "decode": + + def decode_failure(frames): + raise RuntimeError("local decode failed") + + monkeypatch.setattr("transfer_queue.storage.payload_transfer.nixl.decode", decode_failure) + with pytest.raises(RuntimeError, match=f"local {failure_stage} failed"): + asyncio.run( + transfer.get( + control_socket=_ProtocolSocket(), + sender_id="manager", + target_id="storage", + global_indexes=[1], + fields=["value"], + ) + ) + assert len(cancellations) == 1 + assert transfer._failed_get_targets == {"unrelated-target"} + + +class _CapturingSocket: + def __init__(self) -> None: + self.messages: list[Any] = [] + + def send_multipart(self, message: Any, **kwargs: Any) -> None: + self.messages.append(message) + + +def test_deferred_response_copies_identity_and_worker_sends_once(): + future: Future[ZMQMessage] = Future() + response = DeferredResponse(future) + completions: SimpleQueue[tuple[bytes, Future[ZMQMessage]]] = SimpleQueue() + reader, writer = socket.socketpair() + reader.settimeout(1) + poller = zmq.Poller() + poller.register(reader.fileno(), zmq.POLLIN) + shutdown = threading.Event() + identity = bytearray(b"client-a") + worker = _CapturingSocket() + timing_finished: list[bool] = [] + measurement = ExitStack() + measurement.callback(timing_finished.append, True) + measurements = {future: measurement} + try: + _queue_deferred_response(identity, response, completions, writer, shutdown) + identity[:] = b"client-b" + future.set_result( + ZMQMessage.create(request_type=ZMQRequestType.GET_DATA_RESPONSE, sender_id="storage", body={}) + ) + assert dict(poller.poll(1000))[reader.fileno()] == zmq.POLLIN + assert reader.recv(1) == b"\0" + assert timing_finished == [] + _drain_deferred_responses(completions, worker, "storage", measurements) + _drain_deferred_responses(completions, worker, "storage", measurements) + finally: + reader.close() + writer.close() + + assert len(worker.messages) == 1 + assert worker.messages[0][0] == b"client-a" + assert timing_finished == [True] + assert measurements == {} diff --git a/tests/test_payload_transfer.py b/tests/test_payload_transfer.py index 8495b56f..9de86166 100644 --- a/tests/test_payload_transfer.py +++ b/tests/test_payload_transfer.py @@ -29,18 +29,18 @@ parse_payload_transfer_config, ) from transfer_queue.storage.payload_transfer.nixl import PayloadDescriptor -from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import _configure_ucx_environment +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlError, NixlRuntime, _configure_ucx_environment from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType def test_payload_descriptor_preserves_frame_layout(): - descriptor = PayloadDescriptor("framed", 4 + 8 * 2 + 5, (2, 3)) + descriptor = PayloadDescriptor("framed", 5, (2, 3)) descriptor.validate() assert PayloadDescriptor.from_dict(descriptor.to_dict()) == descriptor - with pytest.raises(PayloadTransferError, match="packed payload length"): - PayloadDescriptor("framed", 5, (2, 3)).validate() + with pytest.raises(PayloadTransferError, match="payload length mismatch"): + PayloadDescriptor("framed", 6, (2, 3)).validate() def test_payload_descriptor_requires_frame_layout_and_rejects_negative_lengths(): @@ -94,6 +94,12 @@ def test_nixl_factory_rejects_incomplete_peer_endpoints(): ) +@pytest.mark.parametrize("value", [-1, True, 1.5, "256MB", None]) +def test_nixl_receive_cache_budget_rejects_invalid_values_before_agent_creation(value): + with pytest.raises(NixlError, match="receive_buffer_cache_mb must be a non-negative integer"): + NixlRuntime(receive_buffer_cache_mb=value) + + def test_zmq_payload_transfer_handles_put_and_get_requests(): stored = {} transfer = ZmqPayloadTransfer() diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 53d2b18f..5798156d 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -40,6 +40,8 @@ backend: # To use NIXL-UCX, replace the backend above with nixl-ucx. # Optional UCX environment variables; omit to use UCX defaults. # ucx_env_vars: {} + # Unused NIXL receive buffer cache per runtime in MB (1024^2 bytes); 0 disables caching. + # receive_buffer_cache_mb: 256 # See docs/nixl_ucx_payload.md for UCX transport and device configuration. # SSD offload stores large field values on local SSD to reduce long-lived host-memory use. diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 44ab0d16..679ac95f 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -279,6 +279,29 @@ def _define_controller_metrics(self) -> None: ["storage_unit_id"], registry=r, ) + # Remote runtime snapshots use gauges, like the request snapshots below. + self.storage_payload_metrics = { + name: Gauge(f"tq_storage_payload_{name}", description, ["storage_unit_id"], registry=r) + for name, description in { + "receive_buffer_cache_bytes": "Configured receive buffer cache capacity in bytes", + "registered_bytes": "Registered payload backing bytes", + "idle_receive_bytes": "Cached unused receive bytes", + "pending_receive_bytes": "Prepared receive bytes", + "leased_receive_bytes": "Receive bytes still referenced by data readers", + "quarantined_bytes": "Receive bytes retained after uncertain writes or cleanup failures", + "receive_buffer_registrations": "Cumulative receive buffer registrations", + "receive_buffer_reuses": "Cumulative receive buffer cache hits", + "receive_buffer_evictions": "Cumulative receive buffer deregistrations", + "outstanding_sends": "Queued and running payload sends", + "sender_executors": "Retained sender executors", + "failed_send_peers": "Failed native sender sessions", + "failed_get_targets": "GET targets stopped after uncertain completion", + "queue_wait_seconds": "Cumulative queue wait for started sends in seconds", + "registration_seconds": "Cumulative memory registration time in seconds", + "data_transfer_seconds": "Cumulative native time for writes reaching a terminal status, in seconds", + "total_seconds": "Cumulative started send time including queue wait and cleanup, in seconds", + }.items() + } # ---- Storage request metrics (collected via ZMQ, exposed as gauges) ---- # P50/P99 are pre-computed on the storage unit side and sent via ZMQ, @@ -489,6 +512,16 @@ def collect_storage_metrics(self) -> None: else: gauge.labels(storage_unit_id=label).set(accept_queue.get(key, 0)) + payload_metrics = metrics.get("payload_transfer", {}) + for name, gauge in self.storage_payload_metrics.items(): + if name in payload_metrics: + gauge.labels(storage_unit_id=label).set(payload_metrics[name]) + else: + try: + gauge.remove(label) + except (KeyError, ValueError): + pass + # Per-operation request stats for op_type, op_data in metrics.get("op_stats", {}).items(): self.storage_request_ops.labels(storage_unit_id=label, op_type=op_type).set( diff --git a/transfer_queue/storage/payload_transfer/__init__.py b/transfer_queue/storage/payload_transfer/__init__.py index 621eae0b..f717e268 100644 --- a/transfer_queue/storage/payload_transfer/__init__.py +++ b/transfer_queue/storage/payload_transfer/__init__.py @@ -16,6 +16,7 @@ """Payload transfer strategies for SimpleStorage.""" from transfer_queue.storage.payload_transfer.base import ( + DeferredResponse, PayloadTransfer, PayloadTransferError, PayloadTransferTimeout, @@ -26,6 +27,7 @@ ) __all__ = [ + "DeferredResponse", "PayloadTransfer", "PayloadTransferError", "PayloadTransferTimeout", diff --git a/transfer_queue/storage/payload_transfer/base.py b/transfer_queue/storage/payload_transfer/base.py index 6c9c6d46..2e970253 100644 --- a/transfer_queue/storage/payload_transfer/base.py +++ b/transfer_queue/storage/payload_transfer/base.py @@ -18,6 +18,8 @@ from __future__ import annotations from abc import ABC, abstractmethod +from concurrent.futures import Future +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable if TYPE_CHECKING: @@ -32,7 +34,7 @@ class PayloadTransferError(RuntimeError): """A payload transfer could not be completed safely.""" -class PayloadTransferTimeout(RuntimeError): +class PayloadTransferTimeout(PayloadTransferError): """A payload transfer control message was not answered within the timeout. Distinct from an error the peer reported: only a missing answer is worth retrying on a @@ -45,9 +47,21 @@ def __init__(self, message: str, *, wire_bytes: int = 0): self.wire_bytes = wire_bytes +@dataclass(frozen=True) +class DeferredResponse: + """A response that the SimpleStorage worker must send after completion.""" + + future: Future[ZMQMessage] + + class PayloadTransfer(ABC): """Complete SimpleStorage payload strategy, including its wire protocol.""" + @property + def diagnostics(self) -> dict[str, float | int]: + """Return backend metrics for the storage metrics response.""" + return {} + @abstractmethod async def put( self, @@ -85,7 +99,7 @@ def handle_request( storage_id: str, load_data: Callable[..., dict[str, Any]], store_data: Callable[..., None], - ) -> ZMQMessage | None: + ) -> ZMQMessage | DeferredResponse | None: """Handle a strategy-owned request on a SimpleStorageUnit.""" def bootstrap_info(self) -> dict[str, Any] | None: diff --git a/transfer_queue/storage/payload_transfer/factory.py b/transfer_queue/storage/payload_transfer/factory.py index 28f9ec4a..9458fa54 100644 --- a/transfer_queue/storage/payload_transfer/factory.py +++ b/transfer_queue/storage/payload_transfer/factory.py @@ -16,6 +16,7 @@ """Construction helper for SimpleStorage payload transfer strategies.""" from collections.abc import Mapping +from typing import cast from transfer_queue.storage.payload_transfer.base import PayloadTransfer from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -52,6 +53,7 @@ def create_payload_transfer( return ZmqPayloadTransfer() if backend == "nixl-ucx": from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer + from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB ucx_env_vars = options.get("ucx_env_vars") if ucx_env_vars is not None and not isinstance(ucx_env_vars, Mapping): @@ -60,6 +62,9 @@ def create_payload_transfer( ucx_env_vars=None if ucx_env_vars is None else dict(ucx_env_vars), peer_infos=peer_infos, control_peer_infos=control_peer_infos, + receive_buffer_cache_mb=cast( + int, options.get("receive_buffer_cache_mb", DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB) + ), ) raise RuntimeError(f"unhandled SimpleStorage payload transfer: {backend!r}") diff --git a/transfer_queue/storage/payload_transfer/nixl.py b/transfer_queue/storage/payload_transfer/nixl.py index d4cd070c..41a4897b 100644 --- a/transfer_queue/storage/payload_transfer/nixl.py +++ b/transfer_queue/storage/payload_transfer/nixl.py @@ -28,18 +28,22 @@ import zmq import zmq.asyncio -from transfer_queue.storage.payload_transfer.base import PayloadTransfer, PayloadTransferError, PayloadTransferTimeout -from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlError, NixlRuntime +from transfer_queue.storage.payload_transfer.base import ( + DeferredResponse, + PayloadTransfer, + PayloadTransferError, + PayloadTransferTimeout, +) +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import ( + DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB, + NixlError, + NixlRuntime, + NixlWriteNotStarted, +) from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, unpack_from -from transfer_queue.utils.zmq_utils import ( - ZMQMessage, - ZMQRequestType, - ZMQServerInfo, - create_zmq_socket, - format_zmq_address, -) +from transfer_queue.utils.serial_utils import decode, encode +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) @@ -65,11 +69,10 @@ def validate(self) -> None: raise PayloadTransferError(f"negative payload length for {self.transfer_id}") if any(size < 0 for size in self.frame_sizes): raise PayloadTransferError(f"negative frame length for {self.transfer_id}") - packed_size = 4 + 8 * len(self.frame_sizes) + sum(self.frame_sizes) - if packed_size != self.payload_bytes: + frame_bytes = sum(self.frame_sizes) + if frame_bytes != self.payload_bytes: raise PayloadTransferError( - f"packed payload length mismatch for {self.transfer_id}: " - f"expected {packed_size}, got {self.payload_bytes}" + f"payload length mismatch for {self.transfer_id}: expected {frame_bytes}, got {self.payload_bytes}" ) def to_dict(self) -> dict[str, int | str | list[int]]: @@ -144,31 +147,31 @@ def __init__( ucx_env_vars: dict[str, object] | None = None, peer_infos: Mapping[str, object] | None = None, control_peer_infos: Mapping[str, ZMQServerInfo] | None = None, + *, + receive_buffer_cache_mb: int = DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB, ): self._peer_infos = dict(peer_infos or {}) - self._control_peer_infos = dict(control_peer_infos or {}) - if self._control_peer_infos and set(self._control_peer_infos) != set(self._peer_infos): + if control_peer_infos and set(control_peer_infos) != set(self._peer_infos): raise RuntimeError("SimpleStorage payload transfer endpoints are missing") try: - self._runtime = NixlRuntime(ucx_env_vars) + self._runtime = NixlRuntime(ucx_env_vars, receive_buffer_cache_mb=receive_buffer_cache_mb) except NixlError: raise except Exception as exc: raise NixlError(f"failed to create NIXL runtime: {exc}") from exc self._pending_puts: dict[str, _PendingPut] = {} self._pending_gets: dict[str, _PendingGet] = {} + self._failed_get_targets: set[str] = set() + + @property + def diagnostics(self) -> dict[str, float | int]: + return {**self._runtime.diagnostics, "failed_get_targets": len(self._failed_get_targets)} def bootstrap_info(self) -> dict[str, Any]: return {"endpoint": self.endpoint().to_dict()} def endpoint(self) -> TransferEndpoint: - return TransferEndpoint( - transport=self.transport, - data={ - "agent_name": self._runtime.agent_name, - "agent_metadata": self._runtime.endpoint_metadata(), - }, - ) + return TransferEndpoint(transport=self.transport, data={"agent_name": self._runtime.agent_name}) async def put( self, @@ -181,13 +184,11 @@ async def put( data_parser: Callable[[Any], Any] | None, ) -> int: frames = tuple(encode(data)) - descriptor = PayloadDescriptor( - transfer_id=uuid4().hex, - payload_bytes=calc_packed_size(frames), - frame_sizes=tuple(memoryview(frame).nbytes for frame in frames), - ) - descriptor.validate() - remote_may_be_prepared = True + frame_sizes = tuple(memoryview(frame).nbytes for frame in frames) + descriptor = PayloadDescriptor(transfer_id=uuid4().hex, payload_bytes=sum(frame_sizes), frame_sizes=frame_sizes) + remote_may_be_prepared = False + commit_attempted = False + send_future: Future[None] | None = None try: prepare = ZMQMessage.create( request_type=ZMQRequestType.PUT_DATA_PREPARE, @@ -199,6 +200,7 @@ async def put( "data_parser": data_parser, }, ) + remote_may_be_prepared = True await control_socket.send_multipart(prepare.serialize(), copy=False) ready = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) self._expect(ready, ZMQRequestType.PUT_DATA_READY, target_id) @@ -206,7 +208,8 @@ async def put( raise RuntimeError(f"PUT descriptor changed by storage unit {target_id}") token = ReceiveToken.from_dict(ready.body["receive_token"]) endpoint = self._peer_endpoint(target_id) - await asyncio.wrap_future(self.send(endpoint, token, descriptor, frames)) + send_future = self.send(endpoint, token, descriptor, frames) + await asyncio.wrap_future(send_future) commit = ZMQMessage.create( request_type=ZMQRequestType.PUT_DATA_COMMIT, @@ -214,21 +217,45 @@ async def put( receiver_id=target_id, body={"transfer_id": descriptor.transfer_id}, ) + commit_attempted = True await control_socket.send_multipart(commit.serialize(), copy=False) response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) self._expect(response, ZMQRequestType.PUT_DATA_RESPONSE, target_id) - remote_may_be_prepared = False return descriptor.payload_bytes except BaseException as exc: - if remote_may_be_prepared: - await self._cancel(sender_id, target_id, ZMQRequestType.PUT_DATA_CANCEL, descriptor.transfer_id) + cancel_confirmed = False + if remote_may_be_prepared and not commit_attempted: + safe_to_cancel = True + if send_future is not None and not send_future.cancel(): + # Keep the transaction socket alive until a running WRITE finishes. + try: + await asyncio.wrap_future(send_future) + except NixlWriteNotStarted: + pass + except Exception: + safe_to_cancel = False + if safe_to_cancel: + cancel_confirmed = await self._cancel( + sender_id, target_id, ZMQRequestType.PUT_DATA_CANCEL, descriptor.transfer_id, control_socket + ) if isinstance(exc, zmq.error.Again): - # Only a missing answer is retried by the caller; the transfer_id above is - # fresh per attempt, so replaying the whole put cannot duplicate a commit. + if commit_attempted: + raise PayloadTransferError( + f"NIXL PUT commit outcome is unknown for storage unit {target_id}; automatic retry disabled" + ) from exc + if not cancel_confirmed: + raise PayloadTransferError( + f"NIXL PUT prepare timed out at storage unit {target_id} without confirmed cancellation; " + "automatic retry disabled" + ) from exc raise PayloadTransferTimeout( f"ZMQ recv timeout during put to storage unit {target_id}", wire_bytes=descriptor.payload_bytes, ) from exc + if isinstance(exc, zmq.ZMQError): + raise PayloadTransferError( + f"Error in put to storage unit {target_id}: {type(exc).__name__}: {exc}" + ) from exc raise async def get( @@ -240,10 +267,16 @@ async def get( global_indexes: list[int], fields: list[str], ) -> dict[str, Any]: + if target_id in self._failed_get_targets: + raise NixlError(f"NIXL GET session for {target_id!r} has failed; recreate the payload transfer") + loop = asyncio.get_running_loop() + deadline = loop.time() + TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT transfer_id = uuid4().hex remote_prepared = False receive_prepared = False - descriptor = None + commit_attempted = False + response: ZMQMessage | None = None + response_future: asyncio.Future | None = None try: prepare = ZMQMessage.create( request_type=ZMQRequestType.GET_DATA_PREPARE, @@ -251,13 +284,15 @@ async def get( receiver_id=target_id, body={"global_indexes": global_indexes, "fields": fields, "transfer_id": transfer_id}, ) - await control_socket.send_multipart(prepare.serialize(), copy=False) remote_prepared = True + await control_socket.send_multipart(prepare.serialize(), copy=False) ready = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) self._expect(ready, ZMQRequestType.GET_DATA_READY, target_id) descriptor = PayloadDescriptor.from_dict(ready.body["descriptor"]) if descriptor.transfer_id != transfer_id: raise RuntimeError(f"GET descriptor identity changed by storage unit {target_id}") + if target_id in self._failed_get_targets: + raise NixlError(f"NIXL GET session for {target_id!r} has failed; recreate the payload transfer") token = self.prepare_receive(descriptor) receive_prepared = True commit = ZMQMessage.create( @@ -266,25 +301,73 @@ async def get( receiver_id=target_id, body={ "transfer_id": descriptor.transfer_id, - "receiver_endpoint": self.endpoint().to_dict(), + "receiver_endpoint": TransferEndpoint( + self.transport, {"agent_name": token.data["agent_name"]} + ).to_dict(), "receive_token": token.to_dict(), }, ) + commit_attempted = True await control_socket.send_multipart(commit.serialize(), copy=False) - response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) + response_future = asyncio.ensure_future(control_socket.recv_multipart(copy=False)) + response = ZMQMessage.deserialize(await asyncio.shield(response_future)) self._expect(response, ZMQRequestType.GET_DATA_RESPONSE, target_id) - remote_prepared = False - payload = await asyncio.wrap_future(self.receive(descriptor)) - return decode(unpack_from(payload)) + if response.body.get("transfer_id") != transfer_id: + raise RuntimeError(f"GET completion identity changed by storage unit {target_id}") + frames = self.receive(descriptor) + return decode(list(frames)) except BaseException as exc: - if receive_prepared and descriptor is not None: - self.cancel_receive(descriptor.transfer_id) - if remote_prepared: - await self._cancel(sender_id, target_id, ZMQRequestType.GET_DATA_CANCEL, transfer_id) + cancel_confirmed = False + if receive_prepared: + if commit_attempted and isinstance(exc, asyncio.CancelledError): + # Keep the original receive alive: cancelling the caller does not stop a remote WRITE. + try: + if response_future is None: + response_future = asyncio.ensure_future(control_socket.recv_multipart(copy=False)) + frames = await asyncio.wait_for(asyncio.shield(response_future), max(0, deadline - loop.time())) + response = ZMQMessage.deserialize(list(frames)) + except (Exception, asyncio.CancelledError): + pass + safe_to_release = not commit_attempted or ( + response is not None + and response.sender_id == target_id + and response.body.get("transfer_id") == transfer_id + and ( + response.request_type == ZMQRequestType.GET_DATA_RESPONSE + or ( + response.request_type == ZMQRequestType.PUT_GET_ERROR + and response.body.get("write_not_started") is True + ) + ) + ) + if safe_to_release: + self.cancel_receive(transfer_id) + else: + self._failed_get_targets.add(target_id) + self.quarantine_receive(transfer_id) + if remote_prepared and not commit_attempted: + cancel_confirmed = await self._cancel( + sender_id, target_id, ZMQRequestType.GET_DATA_CANCEL, transfer_id, control_socket + ) if isinstance(exc, zmq.error.Again): - # A retried get prepares a fresh transfer_id, so replaying cannot duplicate. + if commit_attempted: + raise PayloadTransferError( + f"NIXL GET commit outcome is unknown for storage unit {target_id}; automatic retry disabled" + ) from exc + if not cancel_confirmed: + raise PayloadTransferError( + f"NIXL GET prepare timed out at storage unit {target_id} without confirmed cancellation; " + "automatic retry disabled" + ) from exc raise PayloadTransferTimeout(f"ZMQ recv timeout during get from storage unit {target_id}") from exc + if isinstance(exc, zmq.ZMQError): + raise PayloadTransferError( + f"Error getting data from storage unit {target_id}: {type(exc).__name__}: {exc}" + ) from exc raise + finally: + if response_future is not None: + response_future.cancel() def handle_request( self, @@ -293,7 +376,7 @@ def handle_request( storage_id: str, load_data: Callable[..., dict[str, Any]], store_data: Callable[..., None], - ) -> ZMQMessage | None: + ) -> ZMQMessage | DeferredResponse | None: if request.request_type == ZMQRequestType.PUT_DATA_PREPARE: return self._handle_put_prepare(request, storage_id) if request.request_type == ZMQRequestType.PUT_DATA_COMMIT: @@ -336,7 +419,6 @@ def _handle_put_prepare(self, request: ZMQMessage, storage_id: str) -> ZMQMessag def _handle_put_commit(self, request: ZMQMessage, storage_id: str, store_data: Callable[..., None]) -> ZMQMessage: transfer_id = request.body["transfer_id"] - owns_receive = False try: pending = self._pending_puts.get(transfer_id) if pending is None: @@ -344,20 +426,21 @@ def _handle_put_commit(self, request: ZMQMessage, storage_id: str, store_data: C if pending.sender_id != request.sender_id: raise RuntimeError(f"PUT transfer {transfer_id} belongs to another sender") self._pending_puts.pop(transfer_id) - owns_receive = True - payload = self.receive(pending.descriptor).result() + try: + frames = self.receive(pending.descriptor) + except Exception: + self.cancel_receive(transfer_id) + raise with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{storage_id}] PUT commit" ): - store_data(list(pending.global_indexes), decode(unpack_from(payload)), pending.data_parser) + store_data(list(pending.global_indexes), decode(list(frames)), pending.data_parser) return ZMQMessage.create( request_type=ZMQRequestType.PUT_DATA_RESPONSE, sender_id=storage_id, body={"transfer_id": transfer_id}, ) except Exception as exc: - if owns_receive: - self.cancel_receive(transfer_id) return self._error(storage_id, "PUT commit", exc) def _handle_put_cancel(self, request: ZMQMessage, storage_id: str) -> ZMQMessage: @@ -392,12 +475,10 @@ def _handle_get_prepare( target_num_threads=TQ_NUM_THREADS, info=f"[{storage_id}] GET prepare" ): frames = tuple(encode(load_data(request.body["fields"], request.body["global_indexes"]))) + frame_sizes = tuple(memoryview(frame).nbytes for frame in frames) descriptor = PayloadDescriptor( - transfer_id=transfer_id, - payload_bytes=calc_packed_size(frames), - frame_sizes=tuple(memoryview(frame).nbytes for frame in frames), + transfer_id=transfer_id, payload_bytes=sum(frame_sizes), frame_sizes=frame_sizes ) - descriptor.validate() self._pending_gets[transfer_id] = _PendingGet(descriptor, request.sender_id, frames) return ZMQMessage.create( request_type=ZMQRequestType.GET_DATA_READY, @@ -407,8 +488,9 @@ def _handle_get_prepare( except Exception as exc: return self._error(storage_id, "GET prepare", exc) - def _handle_get_commit(self, request: ZMQMessage, storage_id: str) -> ZMQMessage: + def _handle_get_commit(self, request: ZMQMessage, storage_id: str) -> DeferredResponse | ZMQMessage: transfer_id = request.body["transfer_id"] + write_not_started = False try: pending = self._pending_gets.get(transfer_id) if pending is None: @@ -416,16 +498,37 @@ def _handle_get_commit(self, request: ZMQMessage, storage_id: str) -> ZMQMessage if pending.sender_id != request.sender_id: raise RuntimeError(f"GET transfer {transfer_id} belongs to another sender") self._pending_gets.pop(transfer_id) + write_not_started = True endpoint = TransferEndpoint.from_dict(request.body["receiver_endpoint"]) token = ReceiveToken.from_dict(request.body["receive_token"]) - self.send(endpoint, token, pending.descriptor, pending.frames).result() - return ZMQMessage.create( - request_type=ZMQRequestType.GET_DATA_RESPONSE, - sender_id=storage_id, - body={"transfer_id": transfer_id}, - ) + send_future = self.send(endpoint, token, pending.descriptor, pending.frames) + write_not_started = False + response_future: Future[ZMQMessage] = Future() + + def complete_response(completed: Future[None]) -> None: + try: + completed.result() + response = ZMQMessage.create( + request_type=ZMQRequestType.GET_DATA_RESPONSE, + sender_id=storage_id, + body={"transfer_id": transfer_id}, + ) + except Exception as exc: + response = self._error( + storage_id, + "GET commit", + exc, + transfer_id=transfer_id, + write_not_started=completed.cancelled() or isinstance(exc, NixlWriteNotStarted), + ) + response_future.set_result(response) + + send_future.add_done_callback(complete_response) + return DeferredResponse(response_future) except Exception as exc: - return self._error(storage_id, "GET commit", exc) + return self._error( + storage_id, "GET commit", exc, transfer_id=transfer_id, write_not_started=write_not_started + ) def _handle_get_cancel(self, request: ZMQMessage, storage_id: str) -> ZMQMessage: transfer_id = request.body["transfer_id"] @@ -453,16 +556,24 @@ def send( descriptor: PayloadDescriptor, frames: tuple[bytes | bytearray | memoryview, ...] | list[bytes | bytearray | memoryview], ) -> Future[None]: - self._validate_metadata(endpoint, token, descriptor) + if endpoint.transport != self.transport: + raise PayloadTransferError(f"NIXL cannot use endpoint for {endpoint.transport!r}") + if token.data.get("agent_name") != endpoint.data.get("agent_name"): + raise PayloadTransferError("NIXL endpoint and receive token agent names differ") return self._runtime.send(endpoint.data, token.data, descriptor, tuple(frames)) - def receive(self, descriptor: PayloadDescriptor) -> Future[memoryview]: + def receive(self, descriptor: PayloadDescriptor) -> tuple[memoryview, ...]: return self._runtime.receive(descriptor) def cancel_receive(self, transfer_id: str) -> None: self._runtime.cancel_receive(transfer_id) + def quarantine_receive(self, transfer_id: str) -> None: + self._runtime.quarantine_receive(transfer_id) + def close(self) -> None: + self._pending_puts.clear() + self._pending_gets.clear() self._runtime.close() def _peer_endpoint(self, target_id: str) -> TransferEndpoint: @@ -480,45 +591,44 @@ async def _cancel( target_id: str, request_type: ZMQRequestType, transfer_id: str, - ) -> None: - """Send cancellation on a fresh DEALER to isolate late responses.""" - cancel_context = zmq.asyncio.Context() - cancel_socket = None + control_socket: zmq.asyncio.Socket, + ) -> bool: + """Keep PREPARE/CANCEL ordered and confirm cleanup on the transaction's original pipe.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + min(TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, 10) try: - server_info = self._control_peer_infos[target_id] - cancel_socket = create_zmq_socket( - cancel_context, - zmq.DEALER, - server_info.ip, - identity=(f"{sender_id}_cancel_{target_id}_{uuid4().hex[:8]}").encode(), - ) - timeout_ms = min(TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, 10) * 1000 - cancel_socket.setsockopt(zmq.RCVTIMEO, timeout_ms) - cancel_socket.setsockopt(zmq.SNDTIMEO, timeout_ms) - cancel_socket.connect(format_zmq_address(server_info.ip, server_info.ports["put_get_socket"])) cancel = ZMQMessage.create( request_type=request_type, sender_id=sender_id, receiver_id=target_id, body={"transfer_id": transfer_id}, ) - await cancel_socket.send_multipart(cancel.serialize(), copy=False) - response = ZMQMessage.deserialize(await cancel_socket.recv_multipart(copy=False)) + await asyncio.wait_for( + control_socket.send_multipart(cancel.serialize(), copy=False), deadline - loop.time() + ) expected = ( ZMQRequestType.PUT_DATA_RESPONSE if request_type == ZMQRequestType.PUT_DATA_CANCEL else ZMQRequestType.GET_DATA_RESPONSE ) - self._expect(response, expected, target_id) + while True: + # A late READY or PREPARE error can precede the CANCEL response. + frames = await asyncio.wait_for(control_socket.recv_multipart(copy=False), deadline - loop.time()) + response = ZMQMessage.deserialize(frames) + if ( + response.request_type == expected + and response.sender_id == target_id + and response.body.get("transfer_id") == transfer_id + ): + return True except Exception as exc: logger.warning("failed to cancel %s transfer %s at %s: %s", request_type.value, transfer_id, target_id, exc) - finally: - if cancel_socket is not None: - cancel_socket.close(linger=0) - cancel_context.term() + return False @staticmethod def _expect(response: ZMQMessage, expected: ZMQRequestType, target_id: str) -> None: + if response.sender_id != target_id: + raise PayloadTransferError(f"expected response from storage unit {target_id}, got {response.sender_id}") if response.request_type != expected: raise RuntimeError( f"storage unit {target_id} returned {response.request_type}: " @@ -526,21 +636,19 @@ def _expect(response: ZMQMessage, expected: ZMQRequestType, target_id: str) -> N ) @staticmethod - def _error(storage_id: str, operation: str, error: Exception) -> ZMQMessage: + def _error( + storage_id: str, + operation: str, + error: Exception, + *, + transfer_id: str | None = None, + write_not_started: bool = False, + ) -> ZMQMessage: + body: dict[str, Any] = {"message": f"{operation} failed in storage unit {storage_id}: {error}"} + if transfer_id is not None: + body.update(transfer_id=transfer_id, write_not_started=write_not_started) return ZMQMessage.create( request_type=ZMQRequestType.PUT_GET_ERROR, sender_id=storage_id, - body={"message": f"{operation} failed in storage unit {storage_id}: {error}"}, + body=body, ) - - @staticmethod - def _validate_metadata(endpoint: TransferEndpoint, token: ReceiveToken, descriptor: PayloadDescriptor) -> None: - descriptor.validate() - if endpoint.transport != "nixl-ucx": - raise PayloadTransferError(f"NIXL cannot use endpoint for {endpoint.transport!r}") - if token.data.get("agent_name") != endpoint.data.get("agent_name"): - raise PayloadTransferError("NIXL endpoint and receive token agent names differ") - if not isinstance(token.data.get("frame_remote_descs"), bytes): - raise PayloadTransferError("NIXL receive token is missing frame descriptors") - if not isinstance(token.data.get("agent_metadata"), bytes): - raise PayloadTransferError("NIXL receive token is missing agent metadata") diff --git a/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py b/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py index 3bf19482..9b9827d7 100644 --- a/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py +++ b/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Small NIXL runtime used by the SimpleStorage H2H payload adapter.""" +"""NIXL runtime for the SimpleStorage host payload fast path.""" from __future__ import annotations @@ -22,28 +22,42 @@ import socket import threading import time +import weakref from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 +import numpy as np + from transfer_queue.storage.payload_transfer.base import PayloadTransferError from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.serial_utils import initialize_packed_frame_table + +if TYPE_CHECKING: + from transfer_queue.storage.payload_transfer.nixl import PayloadDescriptor logger = get_logger(__name__) DEFAULT_NIXL_TRANSFER_TIMEOUT_SECONDS = 180 +DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB = 256 +# Bound the memory held by small payloads when reusing a larger buffer. +_MAX_RECEIVE_BUFFER_SIZE_RATIO = 2 +_POLL_INTERVAL_SECONDS = 0.0005 class NixlError(PayloadTransferError): """A NIXL runtime operation failed.""" +class NixlWriteNotStarted(NixlError): + """The failed send never exposed the remote receive buffer to a WRITE.""" + + @dataclass class _RegisteredReceiveBuffer: buffer: bytearray - registrations: Any + registration: Any + address: int @property def capacity(self) -> int: @@ -51,26 +65,49 @@ def capacity(self) -> int: @dataclass -class _ReceiveState: - scratch: _RegisteredReceiveBuffer - serialized_frame_descs: bytes - - -@dataclass -class _FrameRegistration: - """One registered source frame kept alive for an active transfer.""" - +class _FrameSource: owner: bytearray | memoryview - registrations: Any + registration: Any | None address: int size: int +@dataclass +class _PeerSender: + executor: ThreadPoolExecutor + outstanding: int = 0 + + def _buffer_address(buffer: bytearray | memoryview) -> int: - """Return the address of a writable, contiguous host buffer.""" - if not buffer: - raise NixlError("NIXL does not support an empty payload buffer") - return ctypes.addressof(ctypes.c_ubyte.from_buffer(buffer)) + """Return the address of a non-empty writable contiguous host buffer.""" + view = memoryview(buffer) + if not view.nbytes: + raise NixlError("NIXL does not support an empty memory region") + if view.readonly or not view.c_contiguous: + raise NixlError("NIXL memory regions must be writable and C-contiguous") + return ctypes.addressof(ctypes.c_ubyte.from_buffer(view)) + + +def _frame_regions(address: int, frame_sizes: tuple[int, ...]) -> list[tuple[int, int, int]]: + """Build non-empty frame regions without narrowing cumulative offsets.""" + regions = [] + offset = 0 + for size in frame_sizes: + if size: + regions.append((address + offset, size, 0)) + offset += size + return regions + + +def _frame_views(owner: Any, frame_sizes: tuple[int, ...]) -> tuple[memoryview, ...]: + """Split a contiguous owner into frame views, preserving empty frames.""" + view = memoryview(owner).cast("B") + frames = [] + offset = 0 + for size in frame_sizes: + frames.append(view[offset : offset + size]) + offset += size + return tuple(frames) def _configure_ucx_environment(ucx_env_vars: dict[str, object] | None) -> dict[str, str]: @@ -94,14 +131,20 @@ def _warn_if_tcp_fallback_possible() -> None: class NixlRuntime: - """Own one NIXL agent and serialize metadata updates safely. + """Own the NIXL agent, registered memory, and transfer completion.""" - SimpleStorage's control plane already orders prepare/send/commit. The - runtime therefore only keeps registered receive buffers alive and waits - for the sender-side NIXL request to finish. - """ - - def __init__(self, ucx_env_vars: dict[str, object] | None = None): + def __init__( + self, + ucx_env_vars: dict[str, object] | None = None, + *, + receive_buffer_cache_mb: int = DEFAULT_NIXL_RECEIVE_BUFFER_CACHE_MB, + ): + if ( + isinstance(receive_buffer_cache_mb, bool) + or not isinstance(receive_buffer_cache_mb, int) + or receive_buffer_cache_mb < 0 + ): + raise NixlError("receive_buffer_cache_mb must be a non-negative integer (MB)") _configure_ucx_environment(ucx_env_vars) try: from nixl import nixl_agent, nixl_agent_config @@ -131,14 +174,36 @@ def __init__(self, ucx_env_vars: dict[str, object] | None = None): os.environ.get("UCX_TLS", "ucx-auto"), ) - self._receives: dict[str, _ReceiveState] = {} - self._reusable_receive_buffer: _RegisteredReceiveBuffer | None = None + self._receives: dict[str, _RegisteredReceiveBuffer | None] = {} + self._idle_receive_buffers: list[_RegisteredReceiveBuffer] = [] + self._leased_receive_buffers: dict[int, _RegisteredReceiveBuffer] = {} + self._quarantined_receive_buffers: list[_RegisteredReceiveBuffer] = [] + self._retained_resources: list[tuple[Any | None, list[_FrameSource]]] = [] + self._registered_sources: dict[int, _FrameSource] = {} self._remote_metadata: dict[str, bytes] = {} - self._deferred_sends: list[tuple[Any, list[_FrameRegistration]]] = [] + self._failed_peers: set[str] = set() + self._peer_senders: dict[str, _PeerSender] = {} self._lock = threading.RLock() + self._closing = threading.Event() self._closed = False self._timeout_seconds = DEFAULT_NIXL_TRANSFER_TIMEOUT_SECONDS - self._send_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="tq-nixl-send") + self._receive_buffer_cache_bytes = receive_buffer_cache_mb * 1024**2 + self._idle_receive_bytes = 0 + self._active_receive_count = 0 + self._active_receive_bytes = 0 + self._receive_cleanup_failed = False + self._receive_working_set_hwm_count = 0 + self._receive_working_set_hwm_bytes = 0 + self._registered_bytes = 0 + self._quarantined_bytes = 0 + self._receive_buffer_acquisitions = 0 + self._receive_buffer_registrations = 0 + self._receive_buffer_reuses = 0 + self._receive_buffer_evictions = 0 + self._registration_seconds = 0.0 + self._queue_wait_seconds = 0.0 + self._data_transfer_seconds = 0.0 + self._total_seconds = 0.0 @staticmethod def _make_agent_name() -> str: @@ -146,140 +211,326 @@ def _make_agent_name() -> str: @property def agent_name(self) -> str: - """Return the stable name advertised by this runtime instance.""" return self._agent_name - def endpoint_metadata(self) -> bytes: - """Return serialized metadata that peers need to address this agent.""" + @property + def diagnostics(self) -> dict[str, float | int]: + """Snapshot payload resources and cumulative transfer costs.""" with self._lock: - self._ensure_open() - return self._agent.get_agent_metadata() - - def prepare_receive(self, descriptor: Any) -> dict[str, Any]: - """Allocate or reuse registered storage for a scatter receive.""" + pending_receive_bytes = sum(scratch.capacity for scratch in self._receives.values() if scratch is not None) + leased_receive_bytes = sum(scratch.capacity for scratch in self._leased_receive_buffers.values()) + return { + "registration_seconds": self._registration_seconds, + "queue_wait_seconds": self._queue_wait_seconds, + "outstanding_sends": sum(sender.outstanding for sender in self._peer_senders.values()), + "sender_executors": len(self._peer_senders), + "failed_send_peers": len(self._failed_peers), + "data_transfer_seconds": self._data_transfer_seconds, + "total_seconds": self._total_seconds, + "registered_bytes": self._registered_bytes, + "quarantined_bytes": self._quarantined_bytes, + "receive_buffer_cache_bytes": self._receive_buffer_cache_bytes, + "idle_receive_bytes": self._idle_receive_bytes, + "pending_receive_bytes": pending_receive_bytes, + "leased_receive_bytes": leased_receive_bytes, + "idle_receive_buffers": len(self._idle_receive_buffers), + "pending_receive_buffers": sum(buffer is not None for buffer in self._receives.values()), + "leased_receive_buffers": len(self._leased_receive_buffers), + "receive_working_set_hwm_count": self._receive_working_set_hwm_count, + "receive_working_set_hwm_bytes": self._receive_working_set_hwm_bytes, + "receive_buffer_acquisitions": self._receive_buffer_acquisitions, + "receive_buffer_registrations": self._receive_buffer_registrations, + "receive_buffer_reuses": self._receive_buffer_reuses, + "receive_buffer_evictions": self._receive_buffer_evictions, + } + + def prepare_receive(self, descriptor: PayloadDescriptor) -> dict[str, Any]: + """Publish frame descriptors and metadata for this receive MR only.""" descriptor.validate() - if not descriptor.frame_sizes or not any(descriptor.frame_sizes): - raise NixlError("NIXL direct-frame receive requires a non-empty frame") with self._lock: self._ensure_open() if descriptor.transfer_id in self._receives: raise NixlError(f"duplicate NIXL receive: {descriptor.transfer_id}") - scratch = self._acquire_receive_buffer(descriptor.payload_bytes) + + scratch = None try: - initialize_packed_frame_table(scratch.buffer, descriptor.frame_sizes) - address = _buffer_address(scratch.buffer) - payload_offset = 4 + 8 * len(descriptor.frame_sizes) - regions = [] - for size in descriptor.frame_sizes: - if size: - regions.append((address + payload_offset, size, 0)) - payload_offset += size - frame_descs = self._agent.get_serialized_descs(self._agent.get_xfer_descs(regions, mem_type="DRAM")) - state = _ReceiveState(scratch, frame_descs) - metadata = self._agent.get_agent_metadata() - self._receives[descriptor.transfer_id] = state + if descriptor.payload_bytes: + scratch = self._acquire_receive_buffer(descriptor.payload_bytes) + regions = _frame_regions(scratch.address, descriptor.frame_sizes) + remote_descs = self._agent.get_xfer_descs(regions, mem_type="DRAM") + serialized = self._agent.get_serialized_descs(remote_descs) + metadata = self._agent.get_partial_agent_metadata( + scratch.registration, inc_conn_info=True, backends=["UCX"] + ) + else: + serialized = b"" + metadata = b"" except Exception as exc: - self._return_receive_buffer(scratch) - raise NixlError(f"failed to register NIXL receive buffer: {exc}") from exc - result = { + if scratch is not None: + self._cache_or_release_receive_buffer(scratch) + raise NixlError(f"failed to prepare NIXL receive buffer: {exc}") from exc + + self._receives[descriptor.transfer_id] = scratch + if scratch is not None: + self._receive_buffer_acquisitions += 1 + self._active_receive_count += 1 + self._active_receive_bytes += scratch.capacity + self._receive_working_set_hwm_count = max( + self._receive_working_set_hwm_count, self._active_receive_count + ) + self._receive_working_set_hwm_bytes = max( + self._receive_working_set_hwm_bytes, self._active_receive_bytes + ) + + return { "agent_name": self._agent_name, "agent_metadata": metadata, - "frame_remote_descs": state.serialized_frame_descs, + "frame_remote_descs": serialized, "payload_bytes": descriptor.payload_bytes, } - return result def send( self, endpoint: dict[str, Any], token: dict[str, Any], - descriptor: Any, + descriptor: PayloadDescriptor, frames: tuple[bytes | bytearray | memoryview, ...], ) -> Future[None]: - """Run the NIXL transfer on the dedicated send thread.""" + """Submit one transfer to the remote peer's single-worker executor.""" + submitted_at = time.monotonic() + deadline = submitted_at + self._timeout_seconds descriptor.validate() - if not descriptor.frame_sizes or not any(descriptor.frame_sizes): - raise NixlError("NIXL direct-frame send requires a non-empty frame") if tuple(memoryview(frame).nbytes for frame in frames) != descriptor.frame_sizes: raise NixlError(f"frame lengths do not match descriptor for {descriptor.transfer_id}") - try: - remote_name, metadata = self._validate_send_metadata(endpoint, token, descriptor) - except Exception as exc: - future: Future[None] = Future() - future.set_exception(exc) - return future + remote_name, metadata = self._validate_send_metadata(endpoint, token, descriptor) with self._lock: self._ensure_open() - return self._send_executor.submit( - self._send_scatter, - remote_name, - metadata, - tuple(frames), - token["frame_remote_descs"], - ) + if remote_name in self._failed_peers: + raise NixlError(f"NIXL peer session {remote_name!r} has failed") + if not descriptor.payload_bytes: + future: Future[None] = Future() + future.set_result(None) + return future + sender = self._peer_senders.get(remote_name) + if sender is None: + idle_peer = next( + (peer for peer, binding in self._peer_senders.items() if binding.outstanding == 0), None + ) + if idle_peer is not None: + sender = self._peer_senders.pop(idle_peer) + else: + sender = _PeerSender(ThreadPoolExecutor(max_workers=1, thread_name_prefix="tq-nixl-send")) + self._peer_senders[remote_name] = sender + sender.outstanding += 1 + try: + future = sender.executor.submit( + self._send_scatter, + remote_name, + metadata, + frames, + token["frame_remote_descs"], + submitted_at, + deadline, + ) + except BaseException: + sender.outstanding -= 1 + if sender.outstanding == 0: + self._peer_senders.pop(remote_name) + # A failed thread start can leave a work item queued. + sender.executor.shutdown(wait=False, cancel_futures=True) + raise + + def retire_send(completed: Future[None]) -> None: + with self._lock: + sender.outstanding -= 1 + + future.add_done_callback(retire_send) + return future + + def receive(self, descriptor: PayloadDescriptor) -> tuple[memoryview, ...]: + """Return zero-copy frame views and lease the receive MR to their owner.""" + descriptor.validate() + with self._lock: + try: + self._ensure_open() + scratch = self._receives[descriptor.transfer_id] + if scratch is None: + self._receives.pop(descriptor.transfer_id) + return tuple(memoryview(b"") for _ in descriptor.frame_sizes) + + owner = np.frombuffer(scratch.buffer, dtype=np.uint8, count=descriptor.payload_bytes) + frames = _frame_views(owner, descriptor.frame_sizes) + key = id(scratch) + weakref.finalize(owner, self._release_lease, key) + self._leased_receive_buffers[key] = scratch + self._receives.pop(descriptor.transfer_id) + return frames + except KeyError as exc: + raise NixlError(f"no prepared NIXL receive for {descriptor.transfer_id}") from exc + except Exception as exc: + raise NixlError(f"failed to expose NIXL receive frames: {exc}") from exc + + def cancel_receive(self, transfer_id: str) -> None: + """Release a receive only when no WRITE can still access it.""" + with self._lock: + if self._closed: + return + scratch = self._receives.pop(transfer_id, None) + if scratch is not None: + self._active_receive_count -= 1 + self._active_receive_bytes -= scratch.capacity + self._cache_or_release_receive_buffer(scratch) + + def quarantine_receive(self, transfer_id: str) -> None: + """Keep an exposed receiver out of the reuse pool until agent teardown.""" + with self._lock: + scratch = self._receives.pop(transfer_id, None) + if scratch is not None: + self._active_receive_count -= 1 + self._active_receive_bytes -= scratch.capacity + self._quarantined_receive_buffers.append(scratch) + self._quarantined_bytes += scratch.capacity + + def close(self) -> None: + """Stop submissions, stop polling, then tear down the agent before owners.""" + with self._lock: + if self._closed: + return + self._closed = True + self._closing.set() + executors = [sender.executor for sender in self._peer_senders.values()] + self._peer_senders.clear() + + for executor in executors: + executor.shutdown(wait=True, cancel_futures=True) + + with self._lock: + agent = self._agent + self._agent = None + retained = self._retained_resources + self._retained_resources = [] + retained_handles = [handle for handle, _ in retained if handle is not None] + retained_sources = [source for _, sources in retained for source in sources] + retained.clear() + + # NIXL handles keep their agent alive. Drop them while every registered + # owner is still retained, then tear down the agent before those owners. + retained_handles.clear() + del agent + + with self._lock: + self._receives.clear() + self._idle_receive_buffers.clear() + self._leased_receive_buffers.clear() + self._quarantined_receive_buffers.clear() + self._registered_sources.clear() + self._remote_metadata.clear() + self._registered_bytes = 0 + self._quarantined_bytes = 0 + self._idle_receive_bytes = 0 + self._active_receive_count = 0 + self._active_receive_bytes = 0 + retained_sources.clear() def _acquire_receive_buffer(self, required_capacity: int) -> _RegisteredReceiveBuffer: - reusable = self._reusable_receive_buffer - if reusable is not None and reusable.capacity >= required_capacity: - self._reusable_receive_buffer = None - return reusable + if self._receive_cleanup_failed: + raise NixlError("NIXL receive memory cleanup failed; recreate the runtime") + candidates = ( + (scratch.capacity, index) + for index, scratch in enumerate(self._idle_receive_buffers) + if required_capacity <= scratch.capacity <= _MAX_RECEIVE_BUFFER_SIZE_RATIO * required_capacity + ) + best = min(candidates, default=None) + if best is not None: + _, index = best + self._receive_buffer_reuses += 1 + scratch = self._idle_receive_buffers.pop(index) + self._idle_receive_bytes -= scratch.capacity + return scratch + + # Make room for a reusable new size without flushing the cache for an + # oversized transfer that cannot be cached when its lease ends. + if required_capacity <= self._receive_buffer_cache_bytes: + while self._idle_receive_bytes > self._receive_buffer_cache_bytes - required_capacity: + oldest = self._idle_receive_buffers.pop(0) + self._idle_receive_bytes -= oldest.capacity + self._deregister_receive_buffer(oldest) + if self._receive_cleanup_failed: + raise NixlError("NIXL receive memory cleanup failed; recreate the runtime") buffer = bytearray(required_capacity) address = _buffer_address(buffer) - registrations = self._agent.register_memory( + started = time.monotonic() + registration = self._agent.register_memory( [(address, required_capacity, 0, "")], mem_type="DRAM", backends=["UCX"] ) - if registrations is None: - raise NixlError("failed to register NIXL receive scratch buffer") - allocated = _RegisteredReceiveBuffer(buffer, registrations) - if reusable is not None: - self._reusable_receive_buffer = None - self._deregister_registration(reusable.registrations) - return allocated - - def _return_receive_buffer(self, scratch: _RegisteredReceiveBuffer) -> None: - reusable = self._reusable_receive_buffer - if reusable is None: - self._reusable_receive_buffer = scratch - elif scratch.capacity > reusable.capacity: - self._deregister_registration(reusable.registrations) - self._reusable_receive_buffer = scratch + self._registration_seconds += time.monotonic() - started + if registration is None: + raise NixlError("failed to register NIXL receive buffer") + self._registered_bytes += required_capacity + self._receive_buffer_registrations += 1 + return _RegisteredReceiveBuffer(buffer, registration, address) + + def _cache_or_release_receive_buffer(self, scratch: _RegisteredReceiveBuffer) -> None: + """Return an unreferenced, safe receiver while the runtime lock is held.""" + if scratch.capacity <= self._receive_buffer_cache_bytes - self._idle_receive_bytes: + self._idle_receive_buffers.append(scratch) + self._idle_receive_bytes += scratch.capacity else: - self._deregister_registration(scratch.registrations) + self._deregister_receive_buffer(scratch) - def _acquire_frame_registration(self, frame: bytes | bytearray | memoryview) -> _FrameRegistration: - view = memoryview(frame) - if view.readonly or not view.c_contiguous: - owner: bytearray | memoryview = bytearray(view) + def _deregister_receive_buffer(self, scratch: _RegisteredReceiveBuffer) -> None: + try: + self._agent.deregister_memory(scratch.registration, backends=["UCX"]) + except Exception as exc: + self._quarantined_receive_buffers.append(scratch) + self._quarantined_bytes += scratch.capacity + self._receive_cleanup_failed = True + logger.warning("failed to deregister NIXL receive buffer: %s", exc) else: - owner = view.cast("B") + self._registered_bytes -= scratch.capacity + self._receive_buffer_evictions += 1 + + def _register_frame_source(self, owner: bytearray | memoryview) -> _FrameSource: + """Register one writable contiguous source while the runtime lock is held.""" address = _buffer_address(owner) - registrations = self._agent.register_memory( - [(address, memoryview(owner).nbytes, 0, "")], mem_type="DRAM", backends=["UCX"] - ) - if registrations is None: + size = memoryview(owner).nbytes + started = time.monotonic() + registration = self._agent.register_memory([(address, size, 0, "")], mem_type="DRAM", backends=["UCX"]) + self._registration_seconds += time.monotonic() - started + if registration is None: raise NixlError("failed to register NIXL source frame") - return _FrameRegistration(owner, registrations, address, memoryview(owner).nbytes) + self._registered_bytes += size + source = _FrameSource(owner, registration, address, size) + self._registered_sources[address] = source + return source - def _release_send_resources(self, handle: Any, registrations: list[_FrameRegistration]) -> bool: - """Release a send only after NIXL confirms its handle can be released.""" - try: - self._agent.release_xfer_handle(handle) - except Exception as exc: - logger.warning("failed to release NIXL transfer handle; keeping source memory registered: %s", exc) - return False - for registration in registrations: - self._deregister_registration(registration.registrations) - return True - - def _reap_deferred_sends(self) -> None: - if not self._deferred_sends: - return - pending = [] - for handle, registrations in self._deferred_sends: - if not self._release_send_resources(handle, registrations): - pending.append((handle, registrations)) - self._deferred_sends = pending + def _acquire_frame_source(self, frame: bytes | bytearray | memoryview) -> _FrameSource: + view = memoryview(frame) + if not view.readonly and view.c_contiguous: + view = view.cast("B") + address = _buffer_address(view) + size = view.nbytes + with self._lock: + self._ensure_open() + for scratch in self._leased_receive_buffers.values(): + if scratch.address <= address and address + size <= scratch.address + scratch.capacity: + return _FrameSource(view, None, address, size) + + # NIXL resolves registrations by address; overlapping sources + # need separate backing memory to release them independently. + overlaps = any( + address < source.address + source.size and source.address < address + size + for source in self._registered_sources.values() + ) + if not overlaps: + return self._register_frame_source(view) + + # Copy outside the lock so unrelated peers can keep progressing. + owner = bytearray(view) + with self._lock: + self._ensure_open() + return self._register_frame_source(owner) def _send_scatter( self, @@ -287,127 +538,164 @@ def _send_scatter( metadata: bytes, frames: tuple[bytes | bytearray | memoryview, ...], serialized_remote_descs: bytes, + submitted_at: float, + deadline: float, ) -> None: - """Send frames directly to matching registered remote regions.""" - frame_registrations: list[_FrameRegistration] = [] + sources: list[_FrameSource] = [] handle = None + error_message = None + peer_state_uncertain = False + write_started = False try: with self._lock: + self._queue_wait_seconds += time.monotonic() - submitted_at self._ensure_open() - self._reap_deferred_sends() - previous = self._remote_metadata.get(remote_name) - if previous != metadata: - if previous is not None: - self._agent.remove_remote_agent(remote_name) - loaded_name = self._agent.add_remote_agent(metadata) - if isinstance(loaded_name, bytes): - loaded_name = loaded_name.decode() - if loaded_name != remote_name: - raise NixlError( - f"NIXL remote agent name mismatch: expected {remote_name!r}, got {loaded_name!r}" - ) - self._remote_metadata[remote_name] = metadata - - for frame in frames: - if memoryview(frame).nbytes: - frame_registrations.append(self._acquire_frame_registration(frame)) + if remote_name in self._failed_peers: + raise NixlError(f"NIXL peer session {remote_name!r} has failed") + + for frame in frames: + self._check_send_deadline(deadline) + if memoryview(frame).nbytes: + sources.append(self._acquire_frame_source(frame)) + + with self._lock: + self._ensure_open() + if remote_name in self._failed_peers: + raise NixlError(f"NIXL peer session {remote_name!r} has failed") local_descs = self._agent.get_xfer_descs( - [(registration.address, registration.size, 0) for registration in frame_registrations], - mem_type="DRAM", + [(source.address, source.size, 0) for source in sources], mem_type="DRAM" ) remote_descs = self._agent.deserialize_descs(serialized_remote_descs) + self._check_send_deadline(deadline) + peer_state_uncertain = True + self._load_remote_metadata(remote_name, metadata) + peer_state_uncertain = False + self._check_send_deadline(deadline) + peer_state_uncertain = True handle = self._agent.initialize_xfer("WRITE", local_descs, remote_descs, remote_name, backends=["UCX"]) + peer_state_uncertain = False + self._check_send_deadline(deadline) + started_transfer = time.monotonic() + peer_state_uncertain = True + write_started = True status = self._agent.transfer(handle) - deadline = time.monotonic() + self._timeout_seconds while status == "PROC": - if time.monotonic() >= deadline: - raise NixlError(f"NIXL WRITE timed out after {self._timeout_seconds:g}s") + if self._closing.is_set(): + raise NixlError("NIXL runtime is closing") + self._check_send_deadline(deadline) with self._lock: - self._ensure_open() status = self._agent.check_xfer_state(handle) if status == "PROC": - time.sleep(0.0005) + time.sleep(_POLL_INTERVAL_SECONDS) + with self._lock: + self._data_transfer_seconds += time.monotonic() - started_transfer if status != "DONE": raise NixlError(f"NIXL WRITE failed with status {status!r}") - except NixlError: - raise except Exception as exc: - raise NixlError(f"NIXL H2H scatter WRITE failed: {exc}") from exc - finally: + error_message = f"NIXL H2H scatter WRITE failed: {exc}" with self._lock: + if handle is not None and not peer_state_uncertain: + # The deadline expired after preparation but before WRITE. + try: + self._agent.release_xfer_handle(handle) + handle = None + except Exception as cleanup_exc: + logger.warning("failed to release unposted NIXL transfer handle: %s", cleanup_exc) if handle is None: - for registration in frame_registrations: - self._deregister_registration(registration.registrations) - elif not self._release_send_resources(handle, frame_registrations): - self._deferred_sends.append((handle, frame_registrations)) + failed_sources = self._cleanup_sources(sources) + if failed_sources: + self._retained_resources.append((None, failed_sources)) + if peer_state_uncertain or failed_sources: + self._failed_peers.add(remote_name) + else: + self._failed_peers.add(remote_name) + self._retained_resources.append((handle, sources)) + self._total_seconds += time.monotonic() - submitted_at + handle = None + sources = [] + failed_sources = [] + + if error_message is not None: + # A retained Future must not keep the native handle and agent alive. + error_type = NixlError if write_started else NixlWriteNotStarted + raise error_type(error_message) + + with self._lock: + cleanup_failed = False + try: + self._agent.release_xfer_handle(handle) + handle = None + except Exception as exc: + logger.warning("failed to release completed NIXL transfer handle: %s", exc) + self._retained_resources.append((handle, sources)) + cleanup_failed = True + if not cleanup_failed: + failed_sources = self._cleanup_sources(sources) + if failed_sources: + self._retained_resources.append((None, failed_sources)) + cleanup_failed = bool(failed_sources) + if cleanup_failed: + self._failed_peers.add(remote_name) + self._total_seconds += time.monotonic() - submitted_at + + def _check_send_deadline(self, deadline: float) -> None: + if time.monotonic() >= deadline: + raise NixlError(f"NIXL send timed out after {self._timeout_seconds:g}s (including queue wait)") + + def _load_remote_metadata(self, remote_name: str, metadata: bytes) -> None: + previous = self._remote_metadata.get(remote_name) + if previous == metadata: + return + if previous is not None: + self._agent.remove_remote_agent(remote_name) + loaded_name = self._agent.add_remote_agent(metadata) + if isinstance(loaded_name, bytes): + loaded_name = loaded_name.decode() + if loaded_name != remote_name: + raise NixlError(f"NIXL remote agent name mismatch: expected {remote_name!r}, got {loaded_name!r}") + self._remote_metadata[remote_name] = metadata + + def _cleanup_sources(self, sources: list[_FrameSource]) -> list[_FrameSource]: + failed = [] + for source in sources: + if source.registration is None: + continue + try: + self._agent.deregister_memory(source.registration, backends=["UCX"]) + self._registered_bytes -= source.size + self._registered_sources.pop(source.address) + except Exception as exc: + logger.warning("failed to deregister NIXL source memory: %s", exc) + failed.append(source) + return failed + + def _release_lease(self, key: int) -> None: + with self._lock: + scratch = self._leased_receive_buffers.get(key) + if scratch is None or self._closed: + return + self._leased_receive_buffers.pop(key) + self._active_receive_count -= 1 + self._active_receive_bytes -= scratch.capacity + self._cache_or_release_receive_buffer(scratch) + + def _ensure_open(self) -> None: + if self._closed: + raise NixlError("NIXL runtime is closed") @staticmethod def _validate_send_metadata( endpoint: dict[str, Any], token: dict[str, Any], - descriptor: Any, + descriptor: PayloadDescriptor, ) -> tuple[str, bytes]: remote_name = str(token.get("agent_name") or endpoint.get("agent_name") or "") - metadata = token.get("agent_metadata") or endpoint.get("agent_metadata") + metadata = token.get("agent_metadata") if not remote_name or not isinstance(metadata, bytes): - raise NixlError("NIXL endpoint is missing remote agent metadata") + raise NixlError("NIXL receive token is missing current agent metadata") if int(token.get("payload_bytes", -1)) != descriptor.payload_bytes: raise NixlError("NIXL receive token length does not match descriptor") if not isinstance(token.get("frame_remote_descs"), bytes): raise NixlError("NIXL receive token is missing frame descriptors") return remote_name, metadata - - def receive(self, descriptor: Any) -> Future[memoryview]: - """Complete a prepared receive and return detached packed payload bytes.""" - descriptor.validate() - future: Future[memoryview] = Future() - with self._lock: - self._ensure_open() - state = self._receives.pop(descriptor.transfer_id, None) - if state is None: - future.set_exception(NixlError(f"no prepared NIXL receive for {descriptor.transfer_id}")) - return future - try: - payload = state.scratch.buffer[: descriptor.payload_bytes] - future.set_result(memoryview(payload)) - except Exception as exc: - future.set_exception(NixlError(f"failed to copy NIXL receive payload: {exc}")) - finally: - self._return_receive_buffer(state.scratch) - return future - - def cancel_receive(self, transfer_id: str) -> None: - """Cancel a prepared receive and deregister its scratch buffer.""" - with self._lock: - state = self._receives.pop(transfer_id, None) - if state is None: - return - self._deregister_registration(state.scratch.registrations) - - def close(self) -> None: - """Stop the send executor and release all NIXL registrations.""" - with self._lock: - if self._closed: - return - self._closed = True - self._send_executor.shutdown(wait=True, cancel_futures=True) - with self._lock: - self._reap_deferred_sends() - for state in self._receives.values(): - self._deregister_registration(state.scratch.registrations) - self._receives.clear() - if self._reusable_receive_buffer is not None: - self._deregister_registration(self._reusable_receive_buffer.registrations) - self._reusable_receive_buffer = None - self._agent = None - - def _ensure_open(self) -> None: - if self._closed: - raise NixlError("NIXL runtime is closed") - - def _deregister_registration(self, registrations: Any) -> None: - try: - self._agent.deregister_memory(registrations, backends=["UCX"]) - except Exception as exc: - logger.warning("failed to deregister NIXL memory: %s", exc) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index f3408f04..0c400b42 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -32,11 +32,14 @@ import os import pickle import shutil +import socket import time import weakref -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path +from queue import Empty, SimpleQueue from threading import Event, Thread from typing import TYPE_CHECKING, Any from uuid import uuid4 @@ -47,7 +50,7 @@ import torch import zmq -from transfer_queue.storage.payload_transfer import PayloadTransfer, create_payload_transfer +from transfer_queue.storage.payload_transfer import DeferredResponse, PayloadTransfer, create_payload_transfer from transfer_queue.storage.payload_transfer.base import KEY_NOT_FOUND_MARKER from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role @@ -125,6 +128,52 @@ class _SSDValueRef: TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0)) +def _queue_deferred_response( + identity: bytes | bytearray | zmq.Frame, + response: DeferredResponse, + completions: SimpleQueue[tuple[bytes, Future[ZMQMessage]]], + wakeup: socket.socket, + shutdown_event: Event, +) -> None: + """Wake the owning worker when an asynchronous response becomes ready.""" + routing_identity = bytes(identity) + + def ready(future: Future[ZMQMessage]) -> None: + if shutdown_event.is_set(): + return + completions.put((routing_identity, future)) + try: + wakeup.send(b"\0") + except (BlockingIOError, OSError): + pass + + response.future.add_done_callback(ready) + + +def _drain_deferred_responses( + completions: SimpleQueue[tuple[bytes, Future[ZMQMessage]]], + worker_socket: zmq.Socket, + storage_id: str, + measurements: dict[Future[ZMQMessage], ExitStack], +) -> None: + """Send completed responses from the only thread that owns the ZMQ socket.""" + while True: + try: + identity, future = completions.get_nowait() + except Empty: + return + with measurements.pop(future): + try: + response = future.result() + except Exception as exc: + response = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, + sender_id=storage_id, + body={"message": f"{storage_id}, deferred response failed: {exc}."}, + ) + worker_socket.send_multipart([identity] + response.serialize(), copy=False) + + class StorageKeyNotFoundError(KeyError): """Raised when a requested global index is absent from a storage unit. @@ -954,6 +1003,13 @@ def _worker_routine(self) -> None: poller = zmq.Poller() poller.register(worker_socket, zmq.POLLIN) + wakeup_reader, wakeup_writer = socket.socketpair() + wakeup_reader.setblocking(False) + wakeup_writer.setblocking(False) + wakeup_fd = wakeup_reader.fileno() + poller.register(wakeup_fd, zmq.POLLIN) + deferred_completions: SimpleQueue[tuple[bytes, Future[ZMQMessage]]] = SimpleQueue() + deferred_measurements: dict[Future[ZMQMessage], ExitStack] = {} logger.info(f"[{self.storage_unit_id}]: worker thread started...") perf_monitor = IntervalPerfMonitor(caller_name=f"{self.storage_unit_id}") @@ -974,13 +1030,43 @@ def _worker_routine(self) -> None: if self._shutdown_event.is_set(): break + if wakeup_fd in socks: + # Wakeup bytes only signal work; the queue owns the responses. + try: + while wakeup_reader.recv(4096): + pass + except BlockingIOError: + pass + try: + _drain_deferred_responses( + deferred_completions, worker_socket, self.storage_unit_id, deferred_measurements + ) + except zmq.ZMQError as exc: + logger.warning(f"[{self.storage_unit_id}]: deferred response send failed: {exc}") + break + if worker_socket in socks: # Keep request handling in a separate scope so payload references are # released before the worker waits for the next request. - self._process_one_worker_request(worker_socket, monitor) + identity, response_msg = self._process_one_worker_request(worker_socket, monitor, deferred_measurements) + if isinstance(response_msg, DeferredResponse): + _queue_deferred_response( + identity, + response_msg, + deferred_completions, + wakeup_writer, + self._shutdown_event, + ) + else: + worker_socket.send_multipart([identity] + response_msg.serialize(), copy=False) logger.info(f"[{self.storage_unit_id}]: worker stopped.") + for measurement in deferred_measurements.values(): + measurement.close() + poller.unregister(wakeup_fd) poller.unregister(worker_socket) + wakeup_reader.close() + wakeup_writer.close() worker_socket.close(linger=0) def _put_decoded_data(self, global_indexes: list[int], field_data: dict[str, Any], data_parser: Any) -> None: @@ -1008,8 +1094,13 @@ def _put_decoded_data(self, global_indexes: list[int], field_data: dict[str, Any ) self.storage_data.put_data(field_data, global_indexes) - def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) -> None: - """Process one storage request and send its response.""" + def _process_one_worker_request( + self, + worker_socket: zmq.Socket, + monitor: Any, + deferred_measurements: dict[Future[ZMQMessage], ExitStack], + ) -> tuple[bytes | zmq.Frame, ZMQMessage | DeferredResponse]: + """Process one storage request and return its response for the worker socket.""" # Messages received from proxy: [identity, serialized_msg_frame1, ...] messages = worker_socket.recv_multipart(copy=False) identity = messages[0] @@ -1029,12 +1120,13 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - sender_id=self.storage_unit_id, body={"message": f"undecodable request: {type(e).__name__}: {e}"}, ) - worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) - return + return identity, error_msg operation = request_msg.request_type started = time.perf_counter() + measurement = ExitStack() + measurement_deferred = False try: # Count each logical PUT/GET once at its first message; NIXL commit/cancel # messages continue that request and must not inflate arrivals. @@ -1053,20 +1145,17 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - ZMQRequestType.GET_DATA_COMMIT: "GET_DATA", }.get(operation) if metric_op is not None: - with monitor.measure(op_type=metric_op): - response_msg = self.payload_transfer.handle_request( - request_msg, - storage_id=self.storage_unit_id, - load_data=self._load_data, - store_data=self._put_decoded_data, - ) - else: - response_msg = self.payload_transfer.handle_request( - request_msg, - storage_id=self.storage_unit_id, - load_data=self._load_data, - store_data=self._put_decoded_data, - ) + measurement.enter_context(monitor.measure(op_type=metric_op)) + response_msg = self.payload_transfer.handle_request( + request_msg, + storage_id=self.storage_unit_id, + load_data=self._load_data, + store_data=self._put_decoded_data, + ) + if isinstance(response_msg, DeferredResponse): + # Keep the request timing open until its asynchronous transfer finishes. + deferred_measurements[response_msg.future] = measurement + measurement_deferred = True if response_msg is None and operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] with monitor.measure(op_type="CLEAR_DATA"): @@ -1099,6 +1188,13 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - }, ) + finally: + if not measurement_deferred: + measurement.close() + + if isinstance(response_msg, DeferredResponse): + return identity, response_msg + response_frames = response_msg.serialize() if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] # This end serializes the get response, so its frames give the true wire size. @@ -1110,7 +1206,7 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - f"samples={len(request_msg.body.get('global_indexes', []))} " f"fields={list(request_msg.body.get('fields', []))}", ) - worker_socket.send_multipart([identity] + response_frames, copy=False) + return identity, response_msg @staticmethod def _field_length(value: Any) -> int | None: @@ -1197,6 +1293,9 @@ def _handle_get_metrics(self) -> ZMQMessage: "non_overflow_drop_delta": stats.non_overflow_drop_delta, "samples": stats.samples, } + payload_metrics = self.payload_transfer.diagnostics + if payload_metrics: + metrics["payload_transfer"] = payload_metrics # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: