diff --git a/tests/test_payload_transfer.py b/tests/test_payload_transfer.py new file mode 100644 index 00000000..2a053470 --- /dev/null +++ b/tests/test_payload_transfer.py @@ -0,0 +1,311 @@ +"""Local ZMQ and simulated NIXL binding checks for the shared request contract.""" + +import asyncio +import ctypes +import json +import sys +import threading +from collections import defaultdict +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +import torch +import zmq.asyncio + +from transfer_queue.storage.payload_transfer import create_payload_transfer +from transfer_queue.storage.payload_transfer.base import ( + Operation, + RequestHeader, + StorageRejected, + TransferError, + WriteOutcomeUnknown, + decode_spans, + encode_spans, +) +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlRuntime, ReadAccessUnknown +from transfer_queue.storage.simple_storage import SimpleStorageUnit +from transfer_queue.utils.serial_utils import decode, unpack_from +from transfer_queue.utils.zmq_utils import STORAGE_MANAGER_IDENTITY_PREFIX, ZMQSocketPool + + +@pytest.fixture +def native_bus(monkeypatch): + """Exercise the real runtime against a binding that checks registrations.""" + agents = {} + bus = SimpleNamespace(lock=threading.Lock(), drop=None, deferred=False, reads=0, read_started=threading.Event()) + + class Agent: + backends = {"UCX": 1} + + def __init__(self, name, config): + self.name = name + self.regions = {} + self.notifications = defaultdict(list) + self.peers = {} + agents[name] = self + + def get_agent_metadata(self): + return json.dumps([self.name, []]).encode() + + def get_partial_agent_metadata(self, registration, **kwargs): + return json.dumps([self.name, registration]).encode() + + def add_remote_agent(self, metadata): + peer, regions = json.loads(metadata) + self.peers.setdefault(peer, {}).update((addr, size) for addr, size, *_ in regions) + return peer + + def remove_remote_agent(self, peer): + self.peers.pop(peer) + + def send_notif(self, peer, packed): + if bus.drop is not None and bus.drop(self.name, peer, packed): + return + with bus.lock: + agents[peer].notifications[self.name].append(packed) + + def get_new_notifs(self, **kwargs): + with bus.lock: + messages = dict(self.notifications) + self.notifications.clear() + return messages + + def register_memory(self, regions, **kwargs): + for addr, size, *_ in regions: + assert addr not in self.regions + self.regions[addr] = size + return regions + + def deregister_memory(self, regions, **kwargs): + for addr, *_ in regions: + self.regions.pop(addr) + + def get_xfer_descs(self, regions, **kwargs): + return regions + + def initialize_xfer(self, operation, local, remote, peer, **kwargs): + assert operation == "READ" + return local[0], remote[0], peer + + def transfer(self, handle): + if bus.deferred: + bus.read_started.set() + return "PROC" + (dst, size, _), (src, remote_size, _), peer = handle + assert remote_size == size <= self.peers[peer][src] + assert size <= self.regions[dst] + assert size <= agents[peer].regions[src] + ctypes.memmove(dst, src, size) + bus.reads += 1 + return "DONE" + + def check_xfer_state(self, handle): + return self.transfer(handle) + + def release_xfer_handle(self, handle): + if bus.deferred: + raise RuntimeError("READ still in flight") + + monkeypatch.setitem(sys.modules, "nixl", SimpleNamespace(nixl_agent=Agent, nixl_agent_config=dict)) + bus.agents = agents + return bus + + +@pytest.fixture +def storage_pair(monkeypatch, native_bus): + monkeypatch.setattr("transfer_queue.storage.simple_storage.get_node_ip_address", lambda: "127.0.0.1") + created = [] + + def create(backend, timeout=3): + config = {"backend": backend} + unit = SimpleStorageUnit.__ray_metadata__.modified_class({"payload_transfer": config}) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}", "put_get_socket", timeout=3) + client = create_payload_transfer(config, owner_id="test-manager", pool=pool, timeout=timeout) + created.append((unit, client, pool, ctx)) + return unit, client + + yield create + for unit, client, pool, ctx in reversed(created): + client.shutdown() + unit.shutdown() + pool.close() + ctx.term() + + +def header(op, **body): + return RequestHeader(op, {"global_indexes": [0], **body}) + + +def update_header(): + return header( + Operation.UPDATE, + merge_fn=lambda old, new: torch.cat([old, new]), + merge_plan={"value": {"held": [True], "dtype": torch.int64, "is_non_tensor": False}}, + ) + + +@pytest.mark.parametrize("backend", ["zmq", "nixl-ucx"]) +def test_payload_contract(storage_pair, native_bus, backend): + unit, client = storage_pair(backend) + if backend == "nixl-ucx": + # The client must not need a SimpleStorage ZMQ payload lease. + assert not hasattr(client, "pool") + + async def run(): + peer = await client.connect(unit.get_payload_transfer_info()) + await client.request( + peer, + header(Operation.PUT, data_parser=lambda data: {"value": [data["value"][0] + 1]}), + send=encode_spans({"value": [torch.arange(8)]}), + ) + reply, _ = await client.request(peer, update_header(), send=encode_spans({"value": [torch.tensor([9])]})) + assert tuple(reply.body["stored_shapes"]["value"]["shapes"][0]) == (9,) + + request = header( + Operation.GET, + fields=["value"], + parser_args=[10], + data_parser=lambda data, args: {"value": [data["value"][0] + args[0]]}, + ) + _, initial = await client.request(peer, request, recv=None) + expected = torch.tensor([11, 12, 13, 14, 15, 16, 17, 18, 19]) + torch.testing.assert_close(decode_spans(initial)["value"][0], expected) + destinations = tuple(memoryview(bytearray(span.nbytes)) for span in initial) + _, result = await client.request(peer, request, recv=destinations) + assert result is destinations + torch.testing.assert_close(decode_spans(result)["value"][0], expected) + + too_short = (memoryview(bytearray(b"x")),) + with pytest.raises(TransferError, match="layout"): + await client.request(peer, request, recv=too_short) + assert bytes(too_short[0]) == b"x" + + with pytest.raises(StorageRejected, match="must not change dict keys"): + await client.request( + peer, header(Operation.GET, fields=["value"], data_parser=lambda *_: {"bad": [1]}), recv=None + ) + results = await asyncio.gather(*(client.request(peer, request, recv=None) for _ in range(5))) + for _, spans in results: + torch.testing.assert_close(decode_spans(spans)["value"][0], expected) + torch.testing.assert_close(decode_spans(initial)["value"][0], expected) + + asyncio.run(run()) + assert unit._arrivals_by_op["UPDATE_DATA"] == 1 + if backend == "nixl-ucx": + assert native_bus.reads == 10 + assert not client._exports + + +def test_nixl_lost_update_reply_holds_source_and_never_replays(storage_pair, native_bus): + unit, client = storage_pair("nixl-ucx", timeout=0.3) + dropped = [] + + def drop(sender, peer, packed): + envelope = decode(unpack_from(packed)) + if envelope["kind"] == "reply" and envelope.get("body", {}).get("stored_shapes"): + dropped.append((sender, peer, packed)) + return True + return False + + native_bus.drop = drop + + async def run(): + peer = await client.connect(unit.get_payload_transfer_info()) + await client.request(peer, header(Operation.PUT), send=encode_spans({"value": [torch.tensor([1])]})) + with pytest.raises(WriteOutcomeUnknown): + await client.request(peer, update_header(), send=encode_spans({"value": [torch.tensor([2])]})) + assert len(client._exports) == 1 + sender, target, packed = dropped[0] + native_bus.drop = None + native_bus.agents[sender].send_notif(target, packed) + _, result = await client.request(peer, header(Operation.GET, fields=["value"]), recv=None) + torch.testing.assert_close(decode_spans(result)["value"][0], torch.tensor([1, 2])) + assert not client._exports + + asyncio.run(run()) + assert unit._arrivals_by_op["UPDATE_DATA"] == 1 + assert len(dropped) == 1 + + +def test_native_read_keeps_registration_when_release_fails(native_bus): + source = NixlRuntime() + destination = NixlRuntime(timeout=0.01) + offer = source.publish(3, lambda buf: buf.__setitem__(slice(0, 3), b"abc")) + try: + native_bus.deferred = True + with pytest.raises(ReadAccessUnknown): + destination.read(source.endpoint(), offer) + assert len(native_bus.agents[destination.name].regions) == 1 + with pytest.raises(TransferError, match="unresolved READ"): + destination.read(source.endpoint(), offer) + finally: + native_bus.deferred = False + destination.close() + source.close() + + +def test_cancelled_get_keeps_offer_until_read_stops(storage_pair, native_bus): + unit, client = storage_pair("nixl-ucx") + + async def run(): + peer = await client.connect(unit.get_payload_transfer_info()) + await client.request(peer, header(Operation.PUT), send=encode_spans({"value": [torch.tensor([1])]})) + native_bus.deferred = True + get = asyncio.create_task(client.request(peer, header(Operation.GET, fields=["value"]), recv=None)) + try: + assert await asyncio.to_thread(native_bus.read_started.wait, 2) + get.cancel() + with pytest.raises(asyncio.CancelledError): + await get + assert len(unit.payload_transfer._exports) == 1 + assert len(client._active_reads) == 1 + finally: + native_bus.deferred = False + deadline = asyncio.get_running_loop().time() + 2 + while unit.payload_transfer._exports: + assert asyncio.get_running_loop().time() < deadline + await asyncio.sleep(0.01) + + asyncio.run(run()) + + +def test_manager_routes_nixl_requests_without_a_zmq_lease(storage_pair, monkeypatch): + from transfer_queue.storage.managers.base import StorageManager + from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager + + unit, _ = storage_pair("nixl-ucx") + context = zmq.asyncio.Context() + + def initialize_base(self, controller_info, config, **kwargs): + self.zmq_context = context + self.storage_manager_id = "test-manager" + + monkeypatch.setattr(StorageManager, "__init__", initialize_base) + monkeypatch.setattr(StorageManager, "close", lambda self: None) + info = unit.get_zmq_server_info() + manager = AsyncSimpleStorageManager( + None, + { + "zmq_info": info, + "payload_transfer": {"backend": "nixl-ucx"}, + "payload_transfer_infos": {info.id: unit.get_payload_transfer_info()}, + }, + ) + + def reject_lease(*args): + pytest.fail("NIXL data request borrowed a ZMQ socket") + + monkeypatch.setattr(manager.storage_rpc_pool, "alease", reject_lease) + + async def run(): + await manager._put_to_single_storage_unit([0], {"value": [torch.tensor([5])]}, info.id) + _, values = await manager._get_from_single_storage_unit([0], ["value"], info.id) + torch.testing.assert_close(values["value"][0], torch.tensor([5])) + + try: + asyncio.run(run()) + finally: + manager.close() + context.term() diff --git a/tests/test_storage_key_race.py b/tests/test_storage_key_race.py index 6d5670bf..54f8e4dc 100644 --- a/tests/test_storage_key_race.py +++ b/tests/test_storage_key_race.py @@ -77,7 +77,11 @@ def test_get_error_reply_is_marked_and_logged_at_debug(storage_data, caplog): ) with caplog.at_level(logging.DEBUG, logger="transfer_queue.storage.simple_storage"): - reply = unit_class._handle_get(unit, request) + from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer + + transport = ZmqPayloadTransfer(owner_id=unit.storage_unit_id) + transport.listen(handle=unit._handle_payload) + reply = transport.handle_legacy(request) assert reply.request_type == ZMQRequestType.GET_ERROR assert KEY_NOT_FOUND_MARKER in reply.body["message"] diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index e0ac6599..24ca935c 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -22,6 +22,9 @@ backend: # SimpleStorage: ZMQ-based in-memory storage for out-of-the-box usage SimpleStorage: + payload_transfer: + backend: zmq # Use nixl-ucx for NIXL notifications and UCX READ. + # ucx_env_vars: {} # Maximum number of experience samples to hold across all storage units. # Set to null for unlimited capacity (no capacity check). total_storage_size: null diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index f0de7a0b..0d6f61c9 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -31,9 +31,11 @@ from typing import Any from uuid import uuid4 +import ray from omegaconf import DictConfig from transfer_queue.storage.bootstrap.provider import StorageBootstrapProvider +from transfer_queue.storage.payload_transfer import payload_backend from transfer_queue.storage.simple_storage import SimpleStorageUnit from transfer_queue.utils.common import get_node_round_robin_scheduling_strategies from transfer_queue.utils.logging_utils import get_logger @@ -50,6 +52,7 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: simple_storage_handles = {} simple_storage_config = dict(conf.backend.SimpleStorage) + transfer_backend = payload_backend(simple_storage_config.get("payload_transfer")) simple_storage_config["_run_id"] = uuid4().hex num_data_storage_units = simple_storage_config["num_data_storage_units"] required_node_resource = simple_storage_config.get("required_node_resource") @@ -85,5 +88,10 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: storage_zmq_info = process_zmq_server_info(simple_storage_handles) backend_name = conf.backend.storage_backend conf.backend[backend_name].zmq_info = storage_zmq_info + if transfer_backend == "nixl-ucx": + endpoints = ray.get([unit.get_payload_transfer_info.remote() for unit in simple_storage_handles.values()]) + conf.backend[backend_name].payload_transfer_infos = { + info.id: endpoint for info, endpoint in zip(storage_zmq_info.values(), endpoints, strict=True) + } return simple_storage_handles diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index d14e67c3..dd6a0719 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -31,10 +31,17 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory -from transfer_queue.storage.simple_storage import ( - KEY_NOT_FOUND_MARKER, - StorageKeyNotFoundError, +from transfer_queue.storage.payload_transfer import ( + Operation, + RequestHeader, + StorageRejected, + TransferTimeout, + create_payload_transfer, + decode_spans, + encode_spans, + payload_backend, ) +from transfer_queue.storage.simple_storage import StorageKeyNotFoundError from transfer_queue.utils.common import log_heavy_operation from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( @@ -43,7 +50,6 @@ ZMQRequestType, ZMQServerInfo, ZMQSocketPool, - frame_nbytes, with_zmq_socket, ) @@ -238,6 +244,21 @@ def __init__( raise ValueError("AsyncSimpleStorageManager requires non-empty 'zmq_info' in config.") self.storage_unit_infos = self._register_servers(server_infos) + transfer_config = config.get("payload_transfer") + if payload_backend(transfer_config) == "zmq": + self._payload_peers = { + unit: {"backend": "zmq", "address": info} for unit, info in self.storage_unit_infos.items() + } + else: + self._payload_peers = dict(config.get("payload_transfer_infos") or {}) + if self._payload_peers.keys() != self.storage_unit_infos.keys(): + raise ValueError("NIXL payload endpoints do not match the SimpleStorage units") + self.payload_transfer = create_payload_transfer( + transfer_config, + owner_id=self.storage_manager_id, + pool=self.storage_rpc_pool, + timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, + ) self._warn_if_pool_can_exhaust_context(len(self.storage_unit_infos)) def _warn_if_pool_can_exhaust_context(self, num_units: int) -> None: @@ -645,64 +666,45 @@ async def update_data( ) return field_schema - @with_storage_unit_socket + async def _request_payload(self, target, operation, body, *, data=None, recv=()): + started = time.perf_counter() + try: + connection = await self.payload_transfer.connect(self._payload_peers[target]) + reply, spans = await self.payload_transfer.request( + connection, RequestHeader(operation, body), send=encode_spans(data) if data is not None else (), + recv=recv, + ) + except TransferTimeout as exc: + verb = operation.name.lower() + size = f"; serialized_mb={exc.wire_bytes / 2**20:.1f}" if operation != Operation.GET else "" + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during {verb} at storage unit " + f"{target} ({self._describe_storage_unit(target)}){size}" + ) from exc + except StorageRejected as exc: + error = StorageKeyNotFoundError if exc.missing_key else RuntimeError + raise error(f"Failed to {operation.name.lower()} data on storage unit {target}: {exc}") from exc + + if data is not None: + log_heavy_operation( + self.storage_manager_id, operation.name.lower(), time.perf_counter() - started, reply.wire_bytes, + f"to {target} at {self._describe_storage_unit(target)} " + f"samples={len(body['global_indexes'])} fields={list(data)}", + ) + return reply, spans + async def _put_to_single_storage_unit( self, global_indexes: list[int], storage_data: dict[str, Any], target_storage_unit: str, data_parser: Callable[[Any], Any] | None = None, - socket: zmq.Socket = None, ): - """ - Send data to a specific storage unit. - """ - - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_DATA, # type: ignore[arg-type] - sender_id=self.storage_manager_id, - receiver_id=target_storage_unit, - body={"global_indexes": global_indexes, "data": storage_data, "data_parser": data_parser}, + await self._request_payload( + target_storage_unit, Operation.PUT, + {"global_indexes": global_indexes, "data_parser": data_parser}, data=storage_data, ) - serialized_bytes = 0 - started = time.perf_counter() - try: - frames = request_msg.serialize() - serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) - await socket.send_multipart(frames, copy=False) - messages = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(messages) - - if response_msg.request_type != ZMQRequestType.PUT_DATA_RESPONSE: - raise RuntimeError( - f"Failed to put data to storage unit {target_storage_unit}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) - # This end serializes the put payload and waits out the round trip, so it holds both - # the true wire size and the end-to-end latency. - log_heavy_operation( - self.storage_manager_id, - "put", - time.perf_counter() - started, - serialized_bytes, - f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " - f"samples={len(global_indexes)} fields={list(storage_data.keys())}", - ) - except zmq.error.Again as e: - raise StorageUnitTimeout( - f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit " - f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; " - f"serialized_mb={serialized_bytes / 2**20:.1f}" - ) from e - except Exception as e: - logger.error( - f"[{self.storage_manager_id}]: Unexpected error during put to storage unit " - f"{target_storage_unit}: {type(e).__name__}: {e}" - ) - raise RuntimeError(f"Error in put to storage unit {target_storage_unit}: {type(e).__name__}: {e}") from e - - @with_storage_unit_socket async def _update_to_single_storage_unit( self, global_indexes: list[int], @@ -710,58 +712,13 @@ async def _update_to_single_storage_unit( merge_fn: Callable[[Any, Any], Any], merge_plan: dict[str, dict[str, Any]], target_storage_unit: str, - socket: zmq.Socket = None, ) -> dict[str, dict[str, Any]]: - """Send an update to one storage unit and return what it stored, per field.""" - body: dict[str, Any] = { - "global_indexes": global_indexes, - "data": storage_data, - "merge_fn": merge_fn, - "merge_plan": merge_plan, - } - - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.UPDATE_DATA, # type: ignore[arg-type] - sender_id=self.storage_manager_id, - receiver_id=target_storage_unit, - body=body, + reply, _ = await self._request_payload( + target_storage_unit, Operation.UPDATE, + {"global_indexes": global_indexes, "merge_fn": merge_fn, "merge_plan": merge_plan}, + data=storage_data, ) - - serialized_bytes = 0 - started = time.perf_counter() - try: - frames = request_msg.serialize() - serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) - await socket.send_multipart(frames, copy=False) - messages = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(messages) - - if response_msg.request_type != ZMQRequestType.UPDATE_DATA_RESPONSE: - raise RuntimeError( - f"Failed to update data on storage unit {target_storage_unit}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) - log_heavy_operation( - self.storage_manager_id, - "update", - time.perf_counter() - started, - serialized_bytes, - f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " - f"samples={len(global_indexes)} fields={list(storage_data)}", - ) - return response_msg.body.get("stored_shapes", {}) - except zmq.error.Again as e: - raise StorageUnitTimeout( - f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during update to storage unit " - f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; " - f"serialized_mb={serialized_bytes / 2**20:.1f}" - ) from e - except Exception as e: - logger.error( - f"[{self.storage_manager_id}]: Unexpected error during update to storage unit " - f"{target_storage_unit}: {type(e).__name__}: {e}" - ) - raise RuntimeError(f"Error in update to storage unit {target_storage_unit}: {type(e).__name__}: {e}") from e + return reply.body.get("stored_shapes", {}) @staticmethod def _pack_field_values(values: list) -> torch.Tensor | NonTensorStack: @@ -926,7 +883,6 @@ async def get_data( return TensorDict(tensor_data, batch_size=len(metadata)) - @with_storage_unit_socket async def _get_from_single_storage_unit( self, global_indexes: list[int], @@ -934,48 +890,16 @@ async def _get_from_single_storage_unit( target_storage_unit: str, data_parser: Callable[[Any, Any], Any] | None = None, parser_args: list | None = None, - socket: zmq.Socket = None, ): """Get data from a single SU by global index keys.""" - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_DATA, # type: ignore[arg-type] - sender_id=self.storage_manager_id, - receiver_id=target_storage_unit, - body={ - "global_indexes": global_indexes, - "fields": fields, - "data_parser": data_parser, - "parser_args": parser_args, - }, + _, spans = await self._request_payload( + target_storage_unit, + Operation.GET, + {"global_indexes": global_indexes, "fields": fields, + "data_parser": data_parser, "parser_args": parser_args}, + recv=None, ) - try: - await socket.send_multipart(request_msg.serialize()) - messages = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(messages) - - if response_msg.request_type == ZMQRequestType.GET_DATA_RESPONSE: - storage_unit_data = response_msg.body["data"] - return fields, storage_unit_data - else: - message = response_msg.body.get("message", "Unknown error") - error_type = StorageKeyNotFoundError if KEY_NOT_FOUND_MARKER in message else RuntimeError - raise error_type(f"Failed to get data from storage unit {target_storage_unit}: {message}") - except zmq.error.Again as e: - raise StorageUnitTimeout( - f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during get from storage unit " - f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}" - ) from e - except StorageKeyNotFoundError: - # Already logged at debug by the storage unit; propagate for the caller to classify. - raise - except Exception as e: - logger.error( - f"[{self.storage_manager_id}]: Unexpected error from storage unit " - f"{target_storage_unit}: {type(e).__name__}: {e}" - ) - raise RuntimeError( - f"Error getting data from storage unit {target_storage_unit}: {type(e).__name__}: {e}" - ) from e + return fields, decode_spans(spans) async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None: """Clear data in remote StorageUnit. @@ -1325,6 +1249,9 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets and context to prevent resource leaks.""" + transfer = getattr(self, "payload_transfer", None) + if transfer is not None: + transfer.shutdown() # Before super(), which may destroy the context these sockets live on. Absent when # the base constructor raised (a failed handshake does), and __del__ still calls this. pool = getattr(self, "storage_rpc_pool", None) diff --git a/transfer_queue/storage/payload_transfer/__init__.py b/transfer_queue/storage/payload_transfer/__init__.py new file mode 100644 index 00000000..1507f3b1 --- /dev/null +++ b/transfer_queue/storage/payload_transfer/__init__.py @@ -0,0 +1,51 @@ +"""Selectable SimpleStorage payload transports.""" + +from collections.abc import Mapping + +from transfer_queue.storage.payload_transfer.base import ( + Operation, + PayloadTransfer, + RequestHeader, + StorageRejected, + TransferTimeout, + decode_spans, + encode_spans, +) + + +def payload_backend(config) -> str: + if config is None: + return "zmq" + if not isinstance(config, Mapping): + raise TypeError("SimpleStorage.payload_transfer must be a mapping") + backend = config.get("backend", "zmq") + if backend not in ("zmq", "nixl-ucx"): + raise ValueError(f"unsupported SimpleStorage payload backend: {backend!r}") + return backend + + +def create_payload_transfer(config, *, owner_id, pool=None, timeout=200) -> PayloadTransfer: + if payload_backend(config) == "zmq": + from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer + + return ZmqPayloadTransfer(owner_id=owner_id, pool=pool) + + from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer + + ucx_env = config.get("ucx_env_vars") + if ucx_env is not None and not isinstance(ucx_env, Mapping): + raise TypeError("SimpleStorage.payload_transfer.ucx_env_vars must be a mapping") + return NixlPayloadTransfer(ucx_env=ucx_env, timeout=timeout) + + +__all__ = [ + "Operation", + "PayloadTransfer", + "RequestHeader", + "StorageRejected", + "TransferTimeout", + "create_payload_transfer", + "decode_spans", + "encode_spans", + "payload_backend", +] diff --git a/transfer_queue/storage/payload_transfer/base.py b/transfer_queue/storage/payload_transfer/base.py new file mode 100644 index 00000000..7bc33fd5 --- /dev/null +++ b/transfer_queue/storage/payload_transfer/base.py @@ -0,0 +1,125 @@ +"""Transport contract for one SimpleStorage payload request.""" + +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable + +from transfer_queue.utils.serial_utils import decode, encode + + +class Operation(str, Enum): + PUT = "PUT_DATA" + GET = "GET_DATA" + UPDATE = "UPDATE_DATA" + + +class TransferError(RuntimeError): + """Transport or protocol failure.""" + + +class TransferTimeout(TransferError): + def __init__(self, message: str, *, wire_bytes: int = 0): + super().__init__(message) + self.wire_bytes = wire_bytes + + +class WriteOutcomeUnknown(TransferError): + """The peer may already have applied the write. Do not replay it.""" + + +class StorageRejected(RuntimeError): + def __init__(self, message: str, *, missing_key: bool = False): + super().__init__(message) + self.missing_key = missing_key + + +@dataclass(frozen=True) +class RequestHeader: + operation: Operation + body: dict[str, Any] + + +@dataclass(frozen=True) +class ReplyHeader: + body: dict[str, Any] + wire_bytes: int = 0 + + +# Each byte view retains its backing object. Spans are ordered codec frames. +Spans = tuple[memoryview, ...] +PeerInfo = dict[str, Any] +StorageHandler = Callable[[RequestHeader, dict[str, Any] | None], tuple[dict[str, Any], dict[str, Any] | None]] + + +def encode_spans(data: dict[str, Any]) -> Spans: + return tuple(memoryview(frame).cast("B") for frame in encode(data)) + + +def decode_spans(spans: Spans) -> dict[str, Any]: + return decode(list(spans)) + + +def validate_destinations(recv: Spans | None) -> None: + if recv is not None and any( + span.readonly or span.format != "B" or span.ndim != 1 or not span.c_contiguous for span in recv + ): + raise TransferError("receive spans must be writable contiguous byte views") + + +def validate_request(header: RequestHeader, send: Spans, recv: Spans | None) -> None: + writing = header.operation in (Operation.PUT, Operation.UPDATE) + if bool(send) != writing: + raise TransferError("PUT and UPDATE require input spans. GET takes no input spans.") + if writing and recv != (): + raise TransferError("PUT and UPDATE expect an acknowledgement without data") + validate_destinations(recv) + + +def place_reply(spans: Spans, recv: Spans | None) -> Spans: + if recv is None: + return spans + validate_destinations(recv) + if tuple(span.nbytes for span in spans) != tuple(span.nbytes for span in recv): + raise TransferError("response span layout does not match the supplied receive spans") + for source, destination in zip(spans, recv, strict=True): + destination[:] = source + return recv + + +class PayloadTransfer(ABC): + """The selected backend owns the logical request through storage acknowledgement.""" + + poll_interval_ms: int | None = None + + def poll_requests(self) -> None: + """Make progress on the unit's existing worker if the backend needs polling.""" + return None + + @abstractmethod + def listen(self, *, handle: StorageHandler, address: Any = None) -> PeerInfo: + """Bind a local storage handler and return the bootstrap information.""" + + @abstractmethod + async def connect(self, peer_info: PeerInfo) -> Any: + """Resolve an opaque connection usable by this backend.""" + + @abstractmethod + async def request( + self, connection: Any, header: RequestHeader, *, send: Spans = (), recv: Spans | None = () + ) -> tuple[ReplyHeader, Spans]: + """Return once response bytes are readable and request sources are reusable. + + None asks the backend to supply response memory. On failure, any memory + still accessible to the transport must remain owned until access stops. + """ + + async def close(self) -> None: + await asyncio.to_thread(self.shutdown) + + @abstractmethod + def shutdown(self) -> None: + """Drain local work and close transport resources for synchronous TQ owners.""" diff --git a/transfer_queue/storage/payload_transfer/nixl.py b/transfer_queue/storage/payload_transfer/nixl.py new file mode 100644 index 00000000..a66b2991 --- /dev/null +++ b/transfer_queue/storage/payload_transfer/nixl.py @@ -0,0 +1,294 @@ +"""NIXL notifications carry requests. Receivers pull published snapshots with READ.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor +from uuid import uuid4 + +from transfer_queue.storage.payload_transfer.base import ( + Operation, + PayloadTransfer, + ReplyHeader, + RequestHeader, + StorageRejected, + TransferError, + TransferTimeout, + WriteOutcomeUnknown, + decode_spans, + encode_spans, + place_reply, + validate_request, +) +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlRuntime, ReadAccessUnknown +from transfer_queue.utils.logging_utils import get_logger +from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, pack_into, unpack_from + +logger = get_logger(__name__) +MAX_CONTROL_BYTES = 1024**2 +PROTOCOL_VERSION = 1 + + +class NixlPayloadTransfer(PayloadTransfer): + poll_interval_ms = 1 + + def __init__(self, *, ucx_env=None, timeout=200): + self._timeout = timeout + self._runtime = NixlRuntime(ucx_env, timeout=min(timeout, 180)) + self._endpoint = self._runtime.endpoint() + self._handle = None + self._lock = threading.RLock() + self._closed = False + self._connections = {} + self._waiters: dict[str, tuple[str, Future]] = {} + self._exports: dict[tuple[str, str], int] = {} + self._active_reads: set[tuple[str, str]] = set() + self._reads = ThreadPoolExecutor(max_workers=1, thread_name_prefix="tq-nixl-read") + + def listen(self, *, handle, address=None): + self._handle = handle + return {"backend": "nixl-ucx", "version": PROTOCOL_VERSION, **self._endpoint} + + async def connect(self, peer_info): + if ( + peer_info.get("backend") != "nixl-ucx" + or peer_info.get("version") != PROTOCOL_VERSION + or not isinstance(peer_info.get("agent"), str) + or not isinstance(peer_info.get("metadata"), bytes) + ): + raise TransferError("invalid NIXL payload peer or protocol version") + with self._lock: + if self._closed: + raise TransferError("NIXL payload transport is closed") + self._runtime.connect(peer_info) + self._connections[peer_info["agent"]] = peer_info + return peer_info + + @staticmethod + def _pack_control(message): + frames = encode({"version": PROTOCOL_VERSION, **message}) + size = calc_packed_size(frames) + if size > MAX_CONTROL_BYTES: + raise TransferError("NIXL control notification exceeds 1 MiB") + packed = bytearray(size) + pack_into(packed, frames) + return bytes(packed) + + def _publish(self, peer, request_id, spans): + key = (peer, request_id) + if key in self._exports: + raise TransferError("duplicate NIXL buffer offer") + size = calc_packed_size(spans) + offer = self._runtime.publish(size, lambda buffer: pack_into(buffer, spans)) + self._exports[key] = offer["address"] + return offer + + def _release_export(self, peer, request_id): + address = self._exports.pop((peer, request_id), None) + if address is not None: + self._runtime.release_export(address) + + def _release_response(self, peer, request_id): + try: + self._runtime.send_notification(peer, self._pack_control({"kind": "release", "id": request_id})) + except Exception: + # The exporter retains the snapshot when release cannot be delivered. + logger.exception("failed to release NIXL response %s", request_id) + + def _read_response(self, peer, request_id, offer): + safe_to_release = True + try: + return tuple(unpack_from(self._runtime.read(peer, offer))) + except ReadAccessUnknown: + safe_to_release = False + raise + finally: + if safe_to_release: + self._release_response(peer, request_id) + with self._lock: + self._active_reads.discard((peer["agent"], request_id)) + + async def request(self, connection, header, *, send=(), recv=()): + validate_request(header, send, recv) + writing = header.operation in (Operation.PUT, Operation.UPDATE) + peer = connection["agent"] + request_id = uuid4().hex + reply = Future() + response = None + with self._lock: + if self._closed: + raise TransferError("NIXL payload transport is closed") + offer = self._publish(peer, request_id, send) if send else None + try: + control = self._pack_control( + { + "kind": "request", + "id": request_id, + "peer": self._endpoint, + "op": header.operation.value, + "body": header.body, + "offer": offer, + } + ) + except BaseException: + self._release_export(peer, request_id) + raise + self._waiters[request_id] = (peer, reply) + + wire_bytes = len(control) + (offer["size"] if offer else 0) + try: + try: + self._runtime.send_notification(connection, control) + deadline = time.monotonic() + self._timeout + while not reply.done(): + self._poll_replies() + if not reply.done(): + if time.monotonic() >= deadline: + raise TransferTimeout("NIXL storage request timed out", wire_bytes=wire_bytes) + await asyncio.sleep(0.001) + response = reply.result() + except Exception as exc: + if writing: + raise WriteOutcomeUnknown(f"NIXL write {request_id} outcome unknown: {exc}") from exc + raise + finally: + with self._lock: + self._waiters.pop(request_id, None) + if response is None and reply.done() and reply.exception() is None and reply.result().get("offer"): + self._release_response(connection, request_id) + + error = response.get("error") + if error is not None: + if response.get("rejected"): + raise StorageRejected(error, missing_key=response.get("missing_key", False)) + if writing and response.get("outcome_unknown"): + raise WriteOutcomeUnknown(error) + raise TransferError(error) + + offer = response.get("offer") + if writing and (not response.get("input_done") or offer is not None): + raise WriteOutcomeUnknown("invalid NIXL write acknowledgement") + if not writing and offer is None: + raise TransferError("NIXL GET response is missing its buffer offer") + if offer is None: + spans = () + else: + with self._lock: + if self._closed: + raise TransferError("NIXL payload transport is closed") + self._active_reads.add((peer, request_id)) + pending = asyncio.wrap_future(self._reads.submit(self._read_response, connection, request_id, offer)) + try: + spans = await asyncio.shield(pending) + except asyncio.CancelledError: + # The READ owns private memory and must finish before releasing the remote offer. + pending.add_done_callback(lambda future: future.exception() if not future.cancelled() else None) + raise + return ReplyHeader(response["body"], wire_bytes), place_reply(spans, recv) + + def _messages(self): + for sender, messages in self._runtime.poll_notifications().items(): + for packed in messages: + try: + if len(packed) > MAX_CONTROL_BYTES: + raise TransferError("oversized NIXL control notification") + message = decode(unpack_from(packed)) + if message["version"] != PROTOCOL_VERSION: + raise TransferError("unsupported NIXL notification protocol") + request_id = message["id"] + if not isinstance(request_id, str) or len(request_id) != 32: + raise TransferError("invalid NIXL request id") + except Exception: + logger.exception("invalid NIXL notification from %s", sender) + continue + yield sender, request_id, message + + def _poll_replies(self): + with self._lock: + if self._closed: + return + for sender, request_id, message in self._messages(): + if message.get("kind") != "reply" or sender not in self._connections: + continue + if message.get("input_done") is True: + self._release_export(sender, request_id) + pending = self._waiters.get(request_id) + if pending is not None and pending[0] == sender: + if not pending[1].done(): + pending[1].set_result(message) + elif message.get("offer") is not None and (sender, request_id) not in self._active_reads: + self._release_response(self._connections[sender], request_id) + + def poll_requests(self): + """Run on the existing storage worker, serial with administrative operations.""" + for sender, request_id, message in self._messages(): + try: + if message.get("kind") == "release": + self._release_export(sender, request_id) + elif message.get("kind") == "request": + self._serve(sender, request_id, message) + except Exception: + logger.exception("NIXL request %s from %s failed", request_id, sender) + + def _serve(self, sender, request_id, message): + peer = message["peer"] + if peer["agent"] != sender: + raise TransferError("NIXL sender does not match its endpoint") + self._runtime.connect(peer) + reply = {"kind": "reply", "id": request_id, "input_done": True} + try: + operation = Operation(message["op"]) + offer = message.get("offer") + if (operation == Operation.GET) != (offer is None): + raise TransferError("invalid NIXL input buffer offer") + data = None + if offer is not None: + reply["input_done"] = False + try: + packed = self._runtime.read(peer, offer) + except ReadAccessUnknown: + raise + except Exception: + reply["input_done"] = True + raise + reply["input_done"] = True + data = decode_spans(tuple(unpack_from(packed))) + result, payload = self._handle(RequestHeader(operation, message["body"]), data) + reply["body"] = result + if payload is not None: + reply["offer"] = self._publish(sender, request_id, encode_spans(payload)) + except StorageRejected as exc: + reply.update(error=str(exc)[:4096], rejected=True, missing_key=exc.missing_key) + except Exception as exc: + logger.exception("NIXL storage request %s failed", request_id) + reply["error"] = str(exc)[:4096] + try: + packed = self._pack_control(reply) + except Exception: + self._release_export(sender, request_id) + # A write may already be accepted even if its result cannot be encoded. + packed = self._pack_control( + { + "kind": "reply", + "id": request_id, + "input_done": reply["input_done"], + "error": "NIXL storage reply cannot be encoded", + "outcome_unknown": message.get("op") != Operation.GET.value, + } + ) + self._runtime.send_notification(peer, packed) + + def shutdown(self): + with self._lock: + if self._closed: + return + self._closed = True + for _, waiter in self._waiters.values(): + if not waiter.done(): + waiter.set_exception(TransferError("NIXL payload transport closed")) + self._reads.shutdown(wait=True) + self._runtime.close() + self._exports.clear() + self._connections.clear() diff --git a/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py b/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py new file mode 100644 index 00000000..19d239fb --- /dev/null +++ b/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py @@ -0,0 +1,207 @@ +"""NIXL UCX agent, registered buffers, and one-sided READ completion.""" + +from __future__ import annotations + +import ctypes +import os +import socket +import threading +import time +from dataclasses import dataclass +from typing import Any +from uuid import uuid4 + +from transfer_queue.storage.payload_transfer.base import TransferError +from transfer_queue.utils.logging_utils import get_logger + +logger = get_logger(__name__) +MAX_REGISTERED_BYTES = 1024**3 +MAX_REGISTERED_REGIONS = 256 +# An agent whose native request cannot be released must retain its memory owners. +_UNDRAINED_AGENTS: list[NixlRuntime] = [] + + +class ReadAccessUnknown(TransferError): + """A READ may still access both buffers.""" + + +@dataclass(eq=False) +class _Region: + buffer: bytearray + address: int + registration: Any + + +class NixlRuntime: + def __init__(self, ucx_env=None, timeout=180): + os.environ.update({str(key): str(value) for key, value in (ucx_env or {}).items()}) + try: + from nixl import nixl_agent, nixl_agent_config + except ImportError as exc: + raise TransferError("install NIXL Python bindings with the UCX plugin") from exc + self.name = f"tq-{socket.gethostname()}-{os.getpid()}-{uuid4().hex[:12]}" + self._agent = nixl_agent( + self.name, + nixl_agent_config(enable_prog_thread=True, backends=["UCX"]), + ) + if "UCX" not in self._agent.backends: + raise TransferError("NIXL UCX plugin is unavailable") + self._timeout = timeout + self._lock = threading.RLock() + self._read_lock = threading.Lock() + self._regions: dict[int, _Region] = {} + self._registered_bytes = 0 + self._peers: set[str] = set() + self._free_exports: list[_Region] = [] + self._unreleased: list[tuple[Any, _Region]] = [] + self._blocked_peers: set[str] = set() + + def endpoint(self): + with self._lock: + return {"agent": self.name, "metadata": self._agent.get_agent_metadata()} + + def connect(self, peer): + with self._lock: + if peer["agent"] not in self._peers: + self._load_peer(peer["agent"], peer["metadata"]) + + def _load_peer(self, name, metadata): + if name in self._blocked_peers: + raise TransferError(f"cannot replace metadata for {name} while a READ may be active") + # Load partial metadata without disconnecting the notification endpoint. + loaded = self._agent.add_remote_agent(metadata) + if isinstance(loaded, bytes): + loaded = loaded.decode() + if loaded != name: + self._agent.remove_remote_agent(loaded) + raise TransferError("NIXL endpoint agent does not match its metadata") + self._peers.add(name) + + def send_notification(self, peer, message): + with self._lock: + self.connect(peer) + # The agent has only UCX. NIXL 1.5's explicit backend argument is incompatible + # with its Python binding, so let the agent select its single plugin. + self._agent.send_notif(peer["agent"], message) + + def poll_notifications(self): + with self._lock: + return self._agent.get_new_notifs(backends=["UCX"]) + + def _allocate(self, size): + if type(size) is not int or not 0 < size <= MAX_REGISTERED_BYTES: + raise TransferError("invalid NIXL buffer size") + if self._registered_bytes + size > MAX_REGISTERED_BYTES or len(self._regions) >= MAX_REGISTERED_REGIONS: + raise TransferError("NIXL registered buffer budget exhausted") + buffer = bytearray(size) + address = ctypes.addressof(ctypes.c_ubyte.from_buffer(buffer)) + registration = self._agent.register_memory([(address, size, 0, "")], mem_type="DRAM", backends=["UCX"]) + if registration is None: + raise TransferError("NIXL memory registration failed") + region = _Region(buffer, address, registration) + self._regions[address] = region + self._registered_bytes += size + return region + + def publish(self, size, fill): + """Publish private bytes whose registration stays live until release.""" + with self._lock: + if type(size) is not int or not 0 < size <= MAX_REGISTERED_BYTES: + raise TransferError("invalid NIXL export size") + region = min( + (region for region in self._free_exports if len(region.buffer) >= size), + key=lambda region: len(region.buffer), + default=None, + ) + if region is None: + region = self._allocate(1 << (size - 1).bit_length()) + else: + self._free_exports.remove(region) + try: + fill(region.buffer) + metadata = self._agent.get_partial_agent_metadata( + region.registration, inc_conn_info=True, backends=["UCX"] + ) + except BaseException: + self.release_export(region.address) + raise + return {"address": region.address, "size": size, "metadata": metadata} + + def release_export(self, address): + with self._lock: + self._free_exports.append(self._regions[address]) + + def release(self, address): + with self._lock: + region = self._regions[address] + try: + self._agent.deregister_memory(region.registration, backends=["UCX"]) + except Exception: + logger.exception("NIXL deregistration failed. Retaining the buffer until agent shutdown.") + return + del self._regions[address] + self._registered_bytes -= len(region.buffer) + + def read(self, peer, offer): + """READ into a private buffer that can outlive the transport after deregistration.""" + size, address = offer["size"], offer["address"] + if type(address) is not int or address <= 0: + raise TransferError("invalid NIXL remote buffer address") + name = peer["agent"] + # Updating a peer's metadata while its READ is active invalidates that request. + with self._read_lock: + with self._lock: + if name in self._blocked_peers: + raise TransferError(f"NIXL peer {name} has an unresolved READ") + self._load_peer(name, offer["metadata"]) + region = self._allocate(size) + handle = None + try: + with self._lock: + local = self._agent.get_xfer_descs([(region.address, size, 0)], mem_type="DRAM") + remote = self._agent.get_xfer_descs([(address, size, 0)], mem_type="DRAM") + handle = self._agent.initialize_xfer("READ", local, remote, name, backends=["UCX"]) + status = self._agent.transfer(handle) + deadline = time.monotonic() + self._timeout + while status == "PROC": + if time.monotonic() >= deadline: + raise TransferError("NIXL READ timed out") + time.sleep(0.001) + with self._lock: + status = self._agent.check_xfer_state(handle) + if status != "DONE": + raise TransferError(f"NIXL READ failed with status {status!r}") + return region.buffer + finally: + with self._lock: + if handle is not None: + try: + self._agent.release_xfer_handle(handle) + except Exception as exc: + self._unreleased.append((handle, region)) + self._blocked_peers.add(name) + raise ReadAccessUnknown(f"NIXL could not release the READ to {name}") from exc + self.release(region.address) + + def close(self): + with self._read_lock, self._lock: + if self._agent is None: + return + remaining = [] + for handle, region in self._unreleased: + try: + self._agent.release_xfer_handle(handle) + except Exception: + remaining.append((handle, region)) + self._unreleased = remaining + if remaining: + if self not in _UNDRAINED_AGENTS: + _UNDRAINED_AGENTS.append(self) + logger.error("NIXL agent still has active READs. Retaining it and its buffers until process exit.") + return + # Keep all exported buffers alive while native teardown revokes their registrations. + self._agent = None + self._free_exports.clear() + self._regions.clear() + self._registered_bytes = 0 + self._peers.clear() diff --git a/transfer_queue/storage/payload_transfer/zmq.py b/transfer_queue/storage/payload_transfer/zmq.py new file mode 100644 index 00000000..d0d2da7b --- /dev/null +++ b/transfer_queue/storage/payload_transfer/zmq.py @@ -0,0 +1,103 @@ +"""Adapter for the existing SimpleStorage multipart protocol.""" + +from __future__ import annotations + +import zmq + +from transfer_queue.storage.payload_transfer.base import ( + Operation, + PayloadTransfer, + ReplyHeader, + RequestHeader, + StorageRejected, + TransferError, + TransferTimeout, + decode_spans, + encode_spans, + place_reply, + validate_request, +) +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, ZMQSocketPool, frame_nbytes + +KEY_NOT_FOUND_MARKER = "TQKeyNotFound" + +REQUEST_TYPES = { + Operation.PUT: ZMQRequestType.PUT_DATA, + Operation.GET: ZMQRequestType.GET_DATA, + Operation.UPDATE: ZMQRequestType.UPDATE_DATA, +} +RESPONSE_TYPES = { + Operation.PUT: ZMQRequestType.PUT_DATA_RESPONSE, + Operation.GET: ZMQRequestType.GET_DATA_RESPONSE, + Operation.UPDATE: ZMQRequestType.UPDATE_DATA_RESPONSE, +} +ERROR_TYPES = { + Operation.PUT: ZMQRequestType.PUT_ERROR, + Operation.GET: ZMQRequestType.GET_ERROR, + Operation.UPDATE: ZMQRequestType.UPDATE_ERROR, +} + + +class ZmqPayloadTransfer(PayloadTransfer): + def __init__(self, *, owner_id: str, pool: ZMQSocketPool | None = None): + self.owner_id = owner_id + self.pool = pool + self._closed = False + + def listen(self, *, handle, address=None): + self._handle = handle + return {"backend": "zmq", "address": address} + + async def connect(self, peer_info): + if peer_info.get("backend") != "zmq" or not isinstance(peer_info.get("address"), ZMQServerInfo): + raise TransferError("invalid ZMQ payload peer") + return peer_info["address"] + + async def request(self, connection, header, *, send=(), recv=()): + if self._closed or self.pool is None: + raise TransferError("ZMQ payload client is unavailable") + validate_request(header, send, recv) + body = dict(header.body) + if send: + body["data"] = decode_spans(send) + request = ZMQMessage.create( + request_type=REQUEST_TYPES[header.operation], + sender_id=self.owner_id, + receiver_id=connection.id, + body=body, + ) + frames = request.serialize() + wire_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) + try: + async with self.pool.alease(connection) as socket: + await socket.send_multipart(frames, copy=False) + reply = ZMQMessage.deserialize(await socket.recv_multipart(copy=False)) + except zmq.error.Again as exc: + raise TransferTimeout("ZMQ storage request timed out", wire_bytes=wire_bytes) from exc + + if reply.request_type != RESPONSE_TYPES[header.operation]: + message = reply.body.get("message", f"unexpected response {reply.request_type}") + raise StorageRejected(message, missing_key=KEY_NOT_FOUND_MARKER in message) + body = dict(reply.body) + spans = encode_spans(body.pop("data")) if "data" in body else () + return ReplyHeader(body, wire_bytes), place_reply(spans, recv) + + def handle_legacy(self, message): + """Translate the legacy wire message at the storage boundary.""" + operation = Operation(message.request_type.name) + body = dict(message.body) + data = body.pop("data", None) + try: + result, payload = self._handle(RequestHeader(operation, body), data) + if payload is not None: + result = {**result, "data": payload} + response_type = RESPONSE_TYPES[operation] + except StorageRejected as exc: + marker = f"[{KEY_NOT_FOUND_MARKER}] " if exc.missing_key else "" + result = {"message": f"{marker}{exc}"} + response_type = ERROR_TYPES[operation] + return ZMQMessage.create(request_type=response_type, sender_id=self.owner_id, body=result) + + def shutdown(self): + # The manager also uses this pool for administrative requests. + self._closed = True diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 8dfa8242..f40de0eb 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -51,6 +51,15 @@ import torch import zmq +from transfer_queue.storage.payload_transfer import ( + Operation, + PayloadTransfer, + RequestHeader, + StorageRejected, + create_payload_transfer, +) +from transfer_queue.storage.payload_transfer.zmq import KEY_NOT_FOUND_MARKER as KEY_NOT_FOUND_MARKER +from transfer_queue.storage.payload_transfer.zmq import REQUEST_TYPES, ZmqPayloadTransfer from transfer_queue.utils import compact_pickle from transfer_queue.utils.common import get_env_bool, limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role @@ -106,7 +115,6 @@ class _SSDValueRef: # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault. -KEY_NOT_FOUND_MARKER = "TQKeyNotFound" # Accept-queue depth for the client-facing ROUTER. A full queue loses connections silently. TQ_STORAGE_ZMQ_BACKLOG = int(os.environ.get("TQ_STORAGE_ZMQ_BACKLOG", 4096)) @@ -972,7 +980,15 @@ def __init__(self, config: dict[str, Any]): self._metrics: TQMetricsExporter | None = None + self.payload_transfer = create_payload_transfer( + config.get("payload_transfer"), owner_id=self.storage_unit_id, + timeout=int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)), + ) self._init_zmq_socket() + unit = weakref.proxy(self) + self._payload_peer = self.payload_transfer.listen( + address=self.zmq_server_info, handle=lambda header, data: unit._dispatch_payload(header, data), + ) self._start_process_put_get() # Register finalizer for graceful cleanup when garbage collected @@ -987,6 +1003,7 @@ def __init__(self, config: dict[str, Any]): self._accept_probe, self.worker_socket, self.storage_data, + self.payload_transfer, ) def shutdown(self) -> None: @@ -1105,6 +1122,7 @@ def _worker_routine(self) -> None: """Worker thread for processing requests.""" logger.info(f"[{self.storage_unit_id}]: worker thread started...") perf_monitor = IntervalPerfMonitor(caller_name=f"{self.storage_unit_id}") + self._payload_monitor = perf_monitor try: self._worker_loop(perf_monitor) @@ -1152,12 +1170,18 @@ def _worker_loop(self, perf_monitor: Any) -> None: """Poll for requests and answer them until shutdown is signalled.""" poller = zmq.Poller() worker_socket = self._connect_worker_socket(poller) + transfer = self.payload_transfer + poll_ms = min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000 + if transfer.poll_interval_ms is not None: + poll_ms = min(poll_ms, transfer.poll_interval_ms) try: while not self._shutdown_event.is_set(): monitor = self._metrics if self._metrics is not None else perf_monitor try: # The event cannot wake a ZMQ poll, so bound idle shutdown latency. - socks = dict(poller.poll(min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000)) + if transfer.poll_interval_ms is not None: + transfer.poll_requests() + socks = dict(poller.poll(poll_ms)) except zmq.error.ContextTerminated: # ZMQ context was terminated, exit gracefully logger.info(f"[{self.storage_unit_id}]: worker stopped gracefully (Context Terminated)") @@ -1211,20 +1235,19 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - started = time.perf_counter() try: - self._requests_arrived += 1 - self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 - + if operation not in REQUEST_TYPES.values(): + self._record_arrival(operation.name) logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") - if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="PUT_DATA"): - response_msg = self._handle_put(request_msg) - elif operation == ZMQRequestType.UPDATE_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="UPDATE_DATA"): - response_msg = self._handle_update(request_msg) - elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="GET_DATA"): - response_msg = self._handle_get(request_msg) + if operation in REQUEST_TYPES.values(): + if isinstance(self.payload_transfer, ZmqPayloadTransfer): + response_msg = self.payload_transfer.handle_legacy(request_msg) + else: + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, + sender_id=self.storage_unit_id, + body={"message": "SimpleStorage payload backend is nixl-ucx"}, + ) elif operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] with monitor.measure(op_type="CLEAR_DATA"): response_msg = self._handle_clear(request_msg) @@ -1290,126 +1313,58 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - ) worker_socket.send_multipart([identity] + response_frames, copy=False) - def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: - """ - Handle put request, add or update data into storage unit. - - Args: - data_parts: ZMQMessage from client. - - Returns: - Put data success response ZMQMessage. - """ - try: - global_indexes = data_parts.body["global_indexes"] - field_data = data_parts.body["data"] # field_data should be a dict. - data_parser = data_parts.body.get("data_parser", None) - - with limit_pytorch_auto_parallel_threads( - target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_put" - ): - if data_parser is not None: - field_data = _apply_data_parser(field_data, data_parser) - self.storage_data.put_data(field_data, global_indexes) - - # After put operation finish, send a message to the client - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_DATA_RESPONSE, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={}, - ) + def _record_arrival(self, operation: str) -> None: + self._requests_arrived += 1 + self._arrivals_by_op[operation] = self._arrivals_by_op.get(operation, 0) + 1 - return response_msg - except Exception as e: - return ZMQMessage.create( - request_type=ZMQRequestType.PUT_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"Failed to put data into storage unit id " - f"#{self.storage_unit_id}, detail error message: {str(e)}" - }, - ) + def _dispatch_payload(self, header: RequestHeader, data: dict[str, Any] | None): + self._record_arrival(header.operation.value) + monitor = self._metrics if self._metrics is not None else self._payload_monitor + with monitor.measure(op_type=header.operation.value): + return self._handle_payload(header, data) - def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage: - """Merge or add fields, write once, and describe what was stored.""" + def _handle_payload(self, header: RequestHeader, data: dict[str, Any] | None): + """Apply storage semantics on the unit worker for every payload backend.""" + body = header.body + operation = header.operation try: - global_indexes = data_parts.body["global_indexes"] - merge_fn = data_parts.body.get("merge_fn") - new_data = data_parts.body.get("data") - merge_plan = data_parts.body.get("merge_plan") - + indexes = body["global_indexes"] with limit_pytorch_auto_parallel_threads( - target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update" + target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] {operation.value}" ): + if operation == Operation.GET: + result = self.storage_data.get_data(body["fields"], indexes) + if (parser := body.get("data_parser")) is not None: + result = _apply_data_parser(result, parser, body.get("parser_args")) + return {}, result + if operation == Operation.PUT: + if (parser := body.get("data_parser")) is not None: + data = _apply_data_parser(data, parser) + self.storage_data.put_data(data, indexes) + return {}, None + + merge_fn = body.get("merge_fn") + merge_plan = body.get("merge_plan") if not callable(merge_fn): raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}") - if not isinstance(new_data, dict) or not new_data: + if not isinstance(data, dict) or not data: raise TypeError("update data must be a non-empty dict") - if not isinstance(merge_plan, dict) or set(merge_plan) != set(new_data): + if not isinstance(merge_plan, dict) or set(merge_plan) != set(data): raise ValueError("merge_plan must describe every update field") - stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, merge_plan) - - return ZMQMessage.create( - request_type=ZMQRequestType.UPDATE_DATA_RESPONSE, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={"stored_shapes": stored_shapes}, - ) - except Exception as e: - return ZMQMessage.create( - request_type=ZMQRequestType.UPDATE_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"Failed to update data in storage unit id " - f"#{self.storage_unit_id}, detail error message: {str(e)}" - }, - ) - - def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage: - """ - Handle get request, return data from storage unit. - - Args: - data_parts: ZMQMessage from client. - - Returns: - Get data success response ZMQMessage, containing target data. - """ - try: - fields = data_parts.body["fields"] - global_indexes = data_parts.body["global_indexes"] - data_parser = data_parts.body.get("data_parser") - - with limit_pytorch_auto_parallel_threads( - target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_get" - ): - result_data = self.storage_data.get_data(fields, global_indexes) - if data_parser is not None: - result_data = _apply_data_parser(result_data, data_parser, data_parts.body.get("parser_args")) - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_DATA_RESPONSE, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "data": result_data, - }, - ) - except Exception as e: - key_not_found = isinstance(e, StorageKeyNotFoundError) - log = logger.debug if key_not_found else logger.error - log( - f"[{self.storage_unit_id}]: _handle_get error, " - f"fields={fields}, global_indexes={global_indexes}: {type(e).__name__}: {e}" - ) - marker = f"[{KEY_NOT_FOUND_MARKER}] " if key_not_found else "" - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"{marker}Failed to get data from storage unit id #{self.storage_unit_id}, " - f"detail error message: {str(e)}" - }, - ) - return response_msg + stored_shapes = self.storage_data.apply_update(indexes, data, merge_fn, merge_plan) + return {"stored_shapes": stored_shapes}, None + except Exception as exc: + missing_key = operation == Operation.GET and isinstance(exc, StorageKeyNotFoundError) + if operation == Operation.GET: + log = logger.debug if missing_key else logger.error + log( + f"[{self.storage_unit_id}]: GET error, fields={body.get('fields')}, " + f"global_indexes={body.get('global_indexes')}: {type(exc).__name__}: {exc}" + ) + raise StorageRejected( + f"Failed to {operation.name.lower()} data on storage unit {self.storage_unit_id}: {exc}", + missing_key=missing_key, + ) from exc def _handle_clear(self, data_parts: ZMQMessage) -> ZMQMessage: """ @@ -1730,6 +1685,7 @@ def _shutdown_resources( accept_probe: "AcceptQueueProbe | None" = None, worker_socket: zmq.Socket | None = None, storage_data: StorageUnitData | None = None, + payload_transfer: PayloadTransfer | None = None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") @@ -1751,6 +1707,8 @@ def _shutdown_resources( worker_socket.close(linger=0) if zmq_context: zmq_context.term() + if payload_transfer is not None: + payload_transfer.shutdown() if storage_data is not None: try: storage_data.close() @@ -1787,3 +1745,6 @@ def get_zmq_server_info(self) -> ZMQServerInfo: ZMQServerInfo containing connection details for this storage unit. """ return self.zmq_server_info + + def get_payload_transfer_info(self) -> dict[str, Any]: + return self._payload_peer