diff --git a/README.md b/README.md index 4df55109..136ab2ff 100644 --- a/README.md +++ b/README.md @@ -72,8 +72,8 @@ In the data plane, we utilize a pluggable design, enabling TransferQueue to inte Specifically, we provide a `StorageManager` abstraction class that defines the core APIs as follows: -- `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None` -- `async def get_data(self, metadata: BatchMeta) -> TensorDict` +- `async def put_data(self, data: TensorDict, metadata: BatchMeta, data_parser=None) -> None` +- `async def get_data(self, metadata: BatchMeta, data_parser=None, parser_args=None) -> TensorDict` - `async def clear_data(self, metadata: BatchMeta, fields: list[str] | None = None) -> None` This class encapsulates the core interaction logic within the TransferQueue system. You only need to write a simple subclass to integrate your custom storage backend. Refer to the [Customize](#customize) section for details. @@ -115,7 +115,7 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag. - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches. -- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields). +- **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields) and, on SimpleStorage, an optional `data_parser` run on the storage units. - **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only. - **(async_)kv_list**: List keys and tags (metadata) in a partition. - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys. @@ -126,6 +126,18 @@ caller stops waiting, so non-idempotent merge operations must not be retried blindly. Updates are atomic within one storage-unit request; a batch spanning multiple units is not a distributed transaction. +**Advanced: parsing on get.** Like `kv_batch_put`, `kv_batch_get` takes a +`data_parser`, here run by each SimpleStorage unit before it replies, so only the +parsed values cross the network (e.g. trimming padding or picking a few frames). +It is called as `data_parser(data, args)`: `data` maps each selected field to the +unit's stored values, one per sample, and `args` holds the matching entries of the +optional per-key `parser_args`. Each field may be parsed differently, but the +result must keep the keys and the number and order of values. `data` holds the +stored objects themselves, so never modify them in place; build new values +instead. As with put, a get with a parser is not retried on a storage-unit +timeout; the call raises, and retrying it yourself may run the parser again, so +keep it free of side effects. See the `kv_batch_get` docstring for details. + **Key Features** - **Redis-style Semantics**: Familiar KV interface (Put/Get/List) for a zero learning curve. diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index fa5d488b..228f22f6 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -727,6 +727,42 @@ def test_kv_batch_get_nonexistent_key(self, controller, tq_api): # Or it might raise an error about keys not found assert "not found" in str(e).lower() or "empty" in str(e).lower() + def test_kv_batch_get_data_parser_parses_each_field_with_per_key_args(self, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("kv_batch_get data_parser is implemented only for SimpleStorage") + + partition_id = "test_partition" + keys = ["parse_0", "parse_1", "parse_2"] + input_ids = torch.tensor([[1, 2, 0, 0], [3, 4, 5, 0], [6, 0, 0, 0]]) + scores = torch.tensor([1.0, 2.0, 3.0]) + tq_api.kv_batch_put( + keys=keys, partition_id=partition_id, fields=TensorDict({"input_ids": input_ids, "scores": scores}, 3) + ) + + def parse(data, args): + data["input_ids"] = [ids[:n] for ids, n in zip(data["input_ids"], args, strict=True)] + data["scores"] = [score * 10 for score in data["scores"]] + return data + + retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id, data_parser=parse, parser_args=[2, 3, 1]) + assert_nested_tensor_equal(retrieved["input_ids"], [input_ids[0, :2], input_ids[1, :3], input_ids[2, :1]]) + assert_tensor_equal(retrieved["scores"], scores * 10) + + retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id) + assert_tensor_equal(retrieved["input_ids"], input_ids) + assert_tensor_equal(retrieved["scores"], scores) + + with pytest.raises(ValueError, match="2 entries for 3 samples"): + tq_api.kv_batch_get(keys=keys, partition_id=partition_id, data_parser=parse, parser_args=[2, 3]) + with pytest.raises(ValueError, match="requires data_parser"): + tq_api.kv_batch_get(keys=keys, partition_id=partition_id, parser_args=[2, 3, 1]) + with pytest.raises(ValueError, match="parser_args needs every key"): + tq_api.kv_batch_get( + keys=["parse_0", "absent"], partition_id=partition_id, data_parser=parse, parser_args=[1, 1] + ) + + tq_api.kv_clear(keys=keys, partition_id=partition_id) + class TestKVBatchGetByMetaE2E: """End-to-end tests for kv_batch_get_by_meta functionality.""" diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index a6875600..609be511 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -286,6 +286,36 @@ async def fake_get(global_indexes, fields, target_storage_unit=None, **kwargs): assert set(called_with["storage_1"]) == {1, 3} +@pytest.mark.asyncio +async def test_get_data_sends_each_unit_the_parser_args_of_its_samples(): + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "test_get_parser" + manager.storage_unit_infos = { + su: ZMQServerInfo(role=Role.STORAGE, id=su, ip="127.0.0.1", ports={"put_get_socket": port}) + for su, port in (("storage_0", 19012), ("storage_1", 19013)) + } + batch_meta = BatchMeta( + global_indexes=[3, 0, 1, 2], + partition_ids=["p0"] * 4, + field_schema={"f": {"dtype": torch.float32, "shape": (2,), "is_nested": False, "is_non_tensor": False}}, + production_status=np.ones(4, dtype=np.int8), + ) + called_with: dict[str, tuple] = {} + + async def fake_get(global_indexes, fields, target_storage_unit=None, data_parser=None, parser_args=None, **_): + called_with[target_storage_unit] = (list(global_indexes), data_parser, parser_args) + return fields, {"f": [torch.zeros(2) for _ in global_indexes]} + + manager._get_from_single_storage_unit = fake_get + + def parser(data, args): + return data + + await manager.get_data(batch_meta, data_parser=parser, parser_args=["a", "b", "c", "d"]) + + assert called_with == {"storage_0": ([0, 2], parser, ["b", "d"]), "storage_1": ([3, 1], parser, ["a", "c"])} + + @pytest.mark.asyncio async def test_clear_data_routes_from_hash(): """clear_data should route using global_idx % num_su (hash routing).""" diff --git a/tests/test_client.py b/tests/test_client.py index aa887656..0fd98761 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -23,6 +23,7 @@ from tensordict import NonTensorStack, TensorDict from transfer_queue import TransferQueueClient +from transfer_queue.client import AsyncTransferQueueClient from transfer_queue.metadata import BatchMeta from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.zmq_utils import ( @@ -1472,3 +1473,27 @@ def test_load_storage_checkpoint_without_manager(self, mock_controller, tmp_path with pytest.raises(RuntimeError, match="Storage manager not initialized"): client.load_storage_checkpoint(ckpt_dir) + + +@pytest.mark.asyncio +async def test_async_get_data_passes_data_parser_only_when_set(): + """Custom managers written as get_data(self, metadata) keep serving plain gets.""" + + class OldStyleManager: + async def get_data(self, metadata): + return TensorDict({"x": torch.zeros(metadata.size)}, batch_size=[metadata.size]) + + client = AsyncTransferQueueClient.__new__(AsyncTransferQueueClient) + client.client_id = "old_manager_client" + client.storage_manager = OldStyleManager() + metadata = BatchMeta( + global_indexes=[0], + partition_ids=["0"], + field_schema={"x": {"dtype": torch.float32, "shape": (1,), "is_nested": False, "is_non_tensor": False}}, + ) + + assert "x" in await client.async_get_data(metadata) + with pytest.raises(TypeError): + await client.async_get_data(metadata, data_parser=lambda data, args: data) + with pytest.raises(ValueError, match="parser_args requires data_parser"): + await client.async_get_data(metadata, parser_args=[1]) diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index dacb0be2..fb2c39bf 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -516,3 +516,15 @@ def test_put_data_custom_backend_meta_length_mismatch_raises_error(test_data_for asyncio.run(manager.put_data(test_data_for_put_data["data"], test_data_for_put_data["metadata"])) assert "does not match" in str(exc_info.value) + + +@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) +def test_get_data_rejects_data_parser(test_data_for_put_data): + mock_storage_client = MagicMock() + config = {"controller_info": MagicMock(), "client_name": "MockClient"} + with patch(f"{STORAGE_CLIENT_FACTORY_PATH}.create", return_value=mock_storage_client): + manager = KVStorageManager(controller_info=MagicMock(), config=config) + + with pytest.raises(NotImplementedError, match="data_parser is not supported"): + asyncio.run(manager.get_data(test_data_for_put_data["metadata"], data_parser=lambda data, args: data)) + mock_storage_client.get.assert_not_called() diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index b8f40a99..001c924e 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -46,8 +46,10 @@ SimpleStorageUnit, SSDFileStore, StorageKeyNotFoundError, + _apply_data_parser, _SSDValueRef, ) +from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder from transfer_queue.utils.zmq_utils import ( STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, @@ -81,11 +83,14 @@ def send_put(self, client_id, global_indexes, field_data, data_parser=None): self.socket.send_multipart(msg.serialize()) return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False)) - def send_get(self, client_id, global_indexes, fields): + def send_get(self, client_id, global_indexes, fields, data_parser=None, parser_args=None): + body = {"global_indexes": global_indexes, "fields": fields} + if data_parser is not None: + body.update(data_parser=data_parser, parser_args=parser_args) msg = ZMQMessage.create( request_type=ZMQRequestType.GET_DATA, sender_id=f"mock_client_{client_id}", - body={"global_indexes": global_indexes, "fields": fields}, + body=body, ) self.socket.send_multipart(msg.serialize()) return ZMQMessage.deserialize(self.socket.recv_multipart(copy=False)) @@ -960,6 +965,59 @@ def __call__(self, field_data): client.close() +def test_storage_unit_get_data_parser_parses_each_field_without_changing_stored_values(storage_setup): + _, put_get_address, storage_ip = storage_setup + client = MockStorageClient(put_get_address, storage_ip) + tokens = torch.arange(8).reshape(2, 4) + images = [["i0", "i1", "i2"], ["j0", "j1"]] + response = client.send_put(0, [0, 1], {"tokens": tokens, "images": images}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE, response.body + + def pick(data, args): + data["tokens"] = [value[:n] for value, (n, _) in zip(data["tokens"], args, strict=True)] + data["images"] = [[value[i] for i in idx] for value, (_, idx) in zip(data["images"], args, strict=True)] + return data + + response = client.send_get(0, [1, 0], ["tokens", "images"], data_parser=pick, parser_args=[(1, [1]), (3, [2, 0])]) + assert response.request_type == ZMQRequestType.GET_DATA_RESPONSE, response.body + torch.testing.assert_close(response.body["data"]["tokens"], [tokens[1, :1], tokens[0, :3]]) + assert response.body["data"]["images"] == [["j1"], ["i2", "i0"]] + + response = client.send_get(0, [0, 1], ["tokens", "images"]) + torch.testing.assert_close(response.body["data"]["tokens"], [tokens[0], tokens[1]]) + assert response.body["data"]["images"] == images + + def drop_one(data, _args): + data["tokens"] = data["tokens"][:-1] + return data + + response = client.send_get(0, [0, 1], ["tokens"], data_parser=drop_one) + assert response.request_type == ZMQRequestType.GET_ERROR + assert "data_parser changed the number of elements" in response.body["message"] + + def unencodable(data, _args): + data["tokens"] = [(x for x in ()) for _ in data["tokens"]] + return data + + response = client.send_get(0, [0, 1], ["tokens"], data_parser=unencodable) + assert response.request_type == ZMQRequestType.GET_ERROR + assert "failed to serialize response" in response.body["message"] + response = client.send_get(0, [0], ["tokens"]) + torch.testing.assert_close(response.body["data"]["tokens"], [tokens[0]]) + + client.close() + + +def test_get_data_parser_view_sends_only_the_parsed_bytes(): + value = torch.arange(12, dtype=torch.int64).view(3, 4) + parsed = _apply_data_parser({"x": [value]}, lambda data, _args: {"x": [v[:2] for v in data["x"]]}, None) + + buffers = MsgpackEncoder().encode(parsed) + assert parsed["x"][0].data_ptr() == value.data_ptr() + assert buffers[1].nbytes == value[:2].nbytes < value.nbytes + assert torch.equal(MsgpackDecoder().decode(buffers)["x"][0], value[:2]) + + def test_storage_unit_data_parser_validation(storage_setup): """Test that invalid data_parser inputs produce clear error messages.""" _, put_get_address, storage_ip = storage_setup diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 70be8985..efab4d60 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -107,6 +107,23 @@ async def test_put_retries_only_without_a_parser(data_parser, expected_attempts) manager.notify_data_update.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data, args: data, 1)]) +async def test_get_retries_only_without_a_parser(data_parser, expected_attempts): + manager = _manager() + manager._get_from_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("no answer")) + _, metadata = _single_sample_batch() + + with ( + patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3), + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + pytest.raises(StorageUnitTimeout), + ): + await manager.get_data(metadata, data_parser=data_parser) + + assert manager._get_from_single_storage_unit.await_count == expected_attempts + + @pytest.mark.asyncio async def test_update_reports_controller_rejection_after_storage_commit(): manager = _manager() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 41ff29d4..7ee7714a 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -522,11 +522,19 @@ async def async_update( field_schema = await self.storage_manager.update_data(metadata, values, merge_fn) return metadata.apply_field_schema(field_schema) - async def async_get_data(self, metadata: BatchMeta) -> TensorDict: + async def async_get_data( + self, + metadata: BatchMeta, + data_parser: Callable[[Any, Any], Any] | None = None, + parser_args: list | None = None, + ) -> TensorDict: """Asynchronously fetch data from storage units and organize into TensorDict. Args: metadata: Batch metadata containing data location information and global indexes + data_parser: Optional ``data_parser(data, args)`` the storage units apply before replying. + See ``transfer_queue.kv_batch_get``. Only supported by SimpleStorage. + parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``. Returns: TensorDict containing: @@ -556,11 +564,22 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: if not metadata or metadata.size == 0 or len(metadata.field_names) == 0: logger.warning(f"[{self.client_id}]: Empty BatchMeta provided to get_data. Returning empty TensorDict.") return TensorDict({}, batch_size=0) + if parser_args is not None: + if data_parser is None: + raise ValueError("parser_args requires data_parser") + if len(parser_args) != metadata.size: + raise ValueError(f"parser_args has {len(parser_args)} entries for {metadata.size} samples") with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_get_data" ): - results = await self.storage_manager.get_data(metadata) + # Passed only when set, so custom managers with get_data(self, metadata) keep working. + if data_parser is None: + results = await self.storage_manager.get_data(metadata) + else: + results = await self.storage_manager.get_data( + metadata, data_parser=data_parser, parser_args=parser_args + ) logger.debug(f"[{self.client_id}]: get_data with {metadata.size} samples successfully.") @@ -1645,11 +1664,19 @@ def update( """ return self._update(metadata, values, merge_fn) - def get_data(self, metadata: BatchMeta) -> TensorDict: + def get_data( + self, + metadata: BatchMeta, + data_parser: Callable[[Any, Any], Any] | None = None, + parser_args: list | None = None, + ) -> TensorDict: """Synchronously fetch data from storage units and organize into TensorDict. Args: metadata: Batch metadata containing data location information and global indexes + data_parser: Optional ``data_parser(data, args)`` the storage units apply before replying. + See ``transfer_queue.kv_batch_get``. Only supported by SimpleStorage. + parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``. Returns: TensorDict containing: @@ -1667,7 +1694,7 @@ def get_data(self, metadata: BatchMeta) -> TensorDict: >>> print(batch) >>> # TensorDict with fields "prompts", "attention_mask", and sample order matching metadata global_indexes """ - return self._get_data(metadata=metadata) + return self._get_data(metadata=metadata, data_parser=data_parser, parser_args=parser_args) def clear_partition(self, partition_id: str): """Synchronously clear the whole partition from all storage units and the controller. diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index b3e0ede4..3d6be893 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -523,7 +523,13 @@ def kv_batch_get_by_meta(meta: KVBatchMeta, select_fields: list[str] | str | Non return tq_client._run_coroutine(async_kv_batch_get_by_meta(meta=meta, select_fields=select_fields)) -def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[str] | str | None = None) -> TensorDict: +def kv_batch_get( + keys: list[str] | str, + partition_id: str, + select_fields: list[str] | str | None = None, + data_parser: Callable[[dict[str, list], list | None], dict[str, list]] | None = None, + parser_args: list | None = None, +) -> TensorDict: """Get data from TransferQueue using user-specified keys. This is a convenience method for retrieving data using keys instead of indexes. @@ -532,6 +538,27 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s keys: Single key or list of keys to retrieve partition_id: Partition containing the keys select_fields: Optional field(s) to retrieve. If None, retrieves all fields + data_parser: Advanced. Optional ``data_parser(data, args) -> dict`` that each storage + unit runs on its share of the batch before replying, so only the parsed values cross + the network (the get-side counterpart of ``kv_batch_put``'s ``data_parser``). + ``data`` maps each selected field to a list with one stored value per sample on that + unit, and ``args`` holds the matching ``parser_args`` entries (None without + ``parser_args``). Each field may be parsed differently, but the returned dict must + keep the keys and the number and order of values per field. Constraints: + + - Do not modify stored values in place (e.g. ``tensor.add_()``, ``value.append()``): + ``data`` holds the stored objects, so the change would persist and also reach + later readers. Build new values instead; replacing entries of ``data`` is safe. + - Keep it free of side effects. Like ``kv_batch_put`` with a parser, a storage-unit + timeout is not retried and the call raises; retrying it yourself may run the parser + again for the same samples. + - It is sent with cloudpickle, so a module-level function must be importable on the + storage units. SSD-offloaded values are read in full before parsing. + + Only supported by SimpleStorage; KV backends raise ``NotImplementedError``. + parser_args: Optional per-key arguments for ``data_parser``, aligned with ``keys``. + Requires ``data_parser``, and every key must exist. They are serialized like + request data, so tuples arrive as lists. Returns: TensorDict with the requested data @@ -539,6 +566,9 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s Raises: ValueError: If keys or partition are not found ValueError: If empty fields exist in any key (sample) + ValueError: If ``parser_args`` is given without ``data_parser``, does not match the + number of keys, or some keys are missing + RuntimeError: If ``data_parser`` fails or breaks its contract on a storage unit Example: >>> import transfer_queue as tq @@ -551,10 +581,28 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s ... partition_id="train", ... select_fields="input_ids" ... ) + >>> # Trim each key's input_ids and pick some of its images on the storage units + >>> def pick(data, args): + ... data["input_ids"] = [ids[:n] for ids, (n, _) in zip(data["input_ids"], args)] + ... data["images"] = [[imgs[i] for i in idx] for imgs, (_, idx) in zip(data["images"], args)] + ... return data + >>> data = tq.kv_batch_get( + ... keys=["sample_1", "sample_2"], + ... partition_id="train", + ... select_fields=["input_ids", "images"], + ... data_parser=pick, + ... parser_args=[(5, [0, 2]), (3, [1])], + ... ) """ tq_client = _maybe_create_tq_client() return tq_client._run_coroutine( - async_kv_batch_get(keys=keys, partition_id=partition_id, select_fields=select_fields) + async_kv_batch_get( + keys=keys, + partition_id=partition_id, + select_fields=select_fields, + data_parser=data_parser, + parser_args=parser_args, + ) ) @@ -938,7 +986,11 @@ async def async_kv_batch_get_by_meta(meta: KVBatchMeta, select_fields: list[str] async def async_kv_batch_get( - keys: list[str] | str, partition_id: str, select_fields: list[str] | str | None = None + keys: list[str] | str, + partition_id: str, + select_fields: list[str] | str | None = None, + data_parser: Callable[[dict[str, list], list | None], dict[str, list]] | None = None, + parser_args: list | None = None, ) -> TensorDict: """Asynchronously get data from TransferQueue using user-specified keys. @@ -948,6 +1000,9 @@ async def async_kv_batch_get( keys: Single key or list of keys to retrieve partition_id: Partition containing the keys select_fields: Optional field(s) to retrieve. If None, retrieves all fields + data_parser: Advanced. Optional parser the storage units run before replying; see + ``kv_batch_get`` for its contract and constraints. + parser_args: Optional per-key arguments for ``data_parser``, aligned with ``keys``. Returns: TensorDict with the requested data @@ -955,6 +1010,9 @@ async def async_kv_batch_get( Raises: ValueError: If keys or partition are not found ValueError: If empty fields exist in any key (sample) + ValueError: If ``parser_args`` is given without ``data_parser``, does not match the + number of keys, or some keys are missing + RuntimeError: If ``data_parser`` fails or breaks its contract on a storage unit Example: >>> import transfer_queue as tq @@ -974,6 +1032,10 @@ async def async_kv_batch_get( if batch_meta.size == 0: raise ValueError("keys or partition were not found!") + # Missing keys are dropped from the meta, which would shift parser_args onto other keys. + num_keys = 1 if isinstance(keys, str) else len(keys) + if parser_args is not None and batch_meta.size != num_keys: + raise ValueError(f"only {batch_meta.size} of {num_keys} keys were found; parser_args needs every key") if select_fields is not None: if isinstance(select_fields, str): @@ -985,7 +1047,7 @@ async def async_kv_batch_get( if not batch_meta.is_ready: raise ValueError("Some fields are not ready in all the requested keys!") - data = await tq_client.async_get_data(batch_meta) + data = await tq_client.async_get_data(batch_meta, data_parser=data_parser, parser_args=parser_args) return data diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 1015172e..b111ac2c 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -334,12 +334,21 @@ async def put_data( raise NotImplementedError("Subclasses must implement put_data") @abstractmethod - async def get_data(self, metadata: BatchMeta) -> TensorDict: + async def get_data( + self, + metadata: BatchMeta, + data_parser: Callable[[Any, Any], Any] | None = None, + parser_args: list | None = None, + ) -> TensorDict: """ Get data from the storage backend. Args: metadata: BatchMeta of the data to be retrieved from the storage. + data_parser: Optional ``data_parser(data, args)`` applied where the data is stored, + before it is sent. See ``transfer_queue.kv_batch_get``. Only supported by + SimpleStorage backend. + parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``. Returns: TensorDict containing the data retrieved from the storage. @@ -823,13 +832,22 @@ async def put_data( per_field_custom_backend_meta, ) - async def get_data(self, metadata: BatchMeta) -> TensorDict: + async def get_data( + self, + metadata: BatchMeta, + data_parser: Callable[[Any, Any], Any] | None = None, + parser_args: list | None = None, + ) -> TensorDict: """ Retrieve tensor data from the backend storage. Fetches tensors using the provided metadata, reconstructs them with the correct shapes and dtypes, and merge them as a TensorDict according to metadata. """ + if data_parser is not None: + raise NotImplementedError( + "data_parser is not supported for KV-based backends (MooncakeStore, Yuanrong, RayStore)." + ) if not metadata.field_names: logger.warning("Attempted to get data, but metadata contains no fields.") return TensorDict({}, batch_size=len(metadata)) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index b5bf88b2..57f22c2d 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -779,7 +779,12 @@ def _pack_field_values(values: list) -> torch.Tensor | NonTensorStack: # Mixed tensor + None — cannot create nested tensor, fall through to NonTensorStack return NonTensorStack(*values) - async def get_data(self, metadata: BatchMeta) -> TensorDict: + async def get_data( + self, + metadata: BatchMeta, + data_parser: Callable[[Any, Any], Any] | None = None, + parser_args: list | None = None, + ) -> TensorDict: """ Retrieve data from remote StorageUnit based on metadata. @@ -787,6 +792,10 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: Args: metadata: BatchMeta that contains metadata for data retrieval. + data_parser: Optional ``data_parser(data, args)`` each SU applies to its part of the + batch before replying. See ``transfer_queue.kv_batch_get``. + parser_args: Optional per-sample arguments aligned with ``metadata.global_indexes``; + each SU receives the slice for its samples as ``args``. Returns: TensorDict containing the retrieved data. @@ -798,6 +807,8 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: return TensorDict({}, batch_size=0) routing = self._group_by_hash(metadata.global_indexes) + # Parser-backed gets are not replayed, matching put: the public API does not constrain parser side effects. + max_attempts = 1 if data_parser is not None else None tasks = [ self._request_with_retry( @@ -809,7 +820,10 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: group.global_indexes, metadata.field_names, target_storage_unit=su_id, + data_parser=data_parser, + parser_args=None if parser_args is None else [parser_args[p] for p in group.batch_positions], ), + max_attempts=max_attempts, ) for su_id, group in routing.items() ] @@ -847,6 +861,8 @@ async def _get_from_single_storage_unit( global_indexes: list[int], fields: list[str], 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.""" @@ -854,7 +870,12 @@ async def _get_from_single_storage_unit( 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}, + body={ + "global_indexes": global_indexes, + "fields": fields, + "data_parser": data_parser, + "parser_args": parser_args, + }, ) try: await socket.send_multipart(request_msg.serialize()) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index ca8ad853..0abc510d 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -832,6 +832,38 @@ def select_dump_schema(schema: dict, source_indexes: list[int], target_indexes: return selected +def _element_count(value: Any) -> int | None: + if hasattr(value, "shape") and isinstance(value.shape, tuple | list) and len(value.shape) > 0: + return value.shape[0] + try: + return len(value) + except Exception: + return None + + +def _apply_data_parser(field_data: dict[str, Any], data_parser: Callable[..., Any], *args: Any) -> dict[str, Any]: + """Run a user ``data_parser(field_data, *args)`` and check it kept every key and element count.""" + if not callable(data_parser): + raise TypeError(f"data_parser must be callable, got {type(data_parser).__name__}") + + original_lengths = {k: _element_count(v) for k, v in field_data.items()} + parsed = data_parser(field_data, *args) + + if not isinstance(parsed, dict): + raise TypeError(f"data_parser must return a dict, got {type(parsed).__name__}") + if set(parsed) != set(original_lengths): + raise ValueError( + f"data_parser must not change dict keys. Original keys: {sorted(original_lengths)}, got: {sorted(parsed)}" + ) + for k, v in parsed.items(): + orig_len, new_len = original_lengths[k], _element_count(v) + if orig_len is not None and new_len is not None and orig_len != new_len: + raise ValueError( + f"data_parser changed the number of elements for key '{k}': expected {orig_len}, got {new_len}" + ) + return parsed + + @ray.remote(num_cpus=1) class SimpleStorageUnit: """A storage unit that provides distributed data storage functionality. @@ -1149,17 +1181,32 @@ def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) - }, ) - 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. - log_heavy_operation( - self.storage_unit_id, - "get", - time.perf_counter() - started, - sum(frame_nbytes(frame) or 0 for frame in response_frames), - f"samples={len(request_msg.body.get('global_indexes', []))} " - f"fields={list(request_msg.body.get('fields', []))}", + # A get data_parser can return a value that fails to encode; reply with an error so the sole worker lives. + try: + response_frames = response_msg.serialize() + except Exception as e: + logger.error( + f"[{self.storage_unit_id}]: failed to serialize response for {operation} " + f"from sender={request_msg.sender_id}: {type(e).__name__}: {e}" ) + response_frames = ZMQMessage.create( + request_type=ZMQRequestType.GET_ERROR # type: ignore[arg-type] + if operation == ZMQRequestType.GET_DATA # type: ignore[arg-type] + else ZMQRequestType.PUT_GET_ERROR, + sender_id=self.storage_unit_id, + body={"message": f"{self.storage_unit_id}: failed to serialize response: {type(e).__name__}: {e}"}, + ).serialize() + else: + if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + # This end serializes the get response, so its frames give the true wire size. + log_heavy_operation( + self.storage_unit_id, + "get", + time.perf_counter() - started, + sum(frame_nbytes(frame) or 0 for frame in response_frames), + 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) def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: @@ -1181,47 +1228,7 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_put" ): if data_parser is not None: - if not callable(data_parser): - raise TypeError(f"data_parser must be callable, got {type(data_parser).__name__}") - - original_keys = set(field_data.keys()) - original_lengths = {} - for k, v in field_data.items(): - if hasattr(v, "shape") and isinstance(v.shape, tuple | list) and len(v.shape) > 0: - original_lengths[k] = v.shape[0] - else: - try: - original_lengths[k] = len(v) - except Exception: - original_lengths[k] = None - - field_data = data_parser(field_data) - - if not isinstance(field_data, dict): - raise TypeError(f"data_parser must return a dict, got {type(field_data).__name__}") - - new_keys = set(field_data.keys()) - if new_keys != original_keys: - raise ValueError( - f"data_parser must not change dict keys. " - f"Original keys: {sorted(original_keys)}, got: {sorted(new_keys)}" - ) - - for k, v in field_data.items(): - if hasattr(v, "shape") and isinstance(v.shape, tuple | list) and len(v.shape) > 0: - new_len = v.shape[0] - else: - try: - new_len = len(v) - except Exception: - new_len = None - - orig_len = original_lengths[k] - if orig_len is not None and new_len is not None and orig_len != new_len: - raise ValueError( - f"data_parser changed the number of elements for key '{k}': " - f"expected {orig_len}, got {new_len}" - ) + 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 @@ -1289,11 +1296,14 @@ def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage: 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]