From 843f6bea6ced6a50c0ff9ccce1813a5fd8b51ef5 Mon Sep 17 00:00:00 2001 From: ji-huazhong Date: Sat, 8 Aug 2026 19:47:21 +0800 Subject: [PATCH] refactor: consolidate KV and controller RPC paths Signed-off-by: ji-huazhong --- AGENTS.md | 6 +- tests/e2e/test_kv_interface_e2e.py | 55 ++++ transfer_queue/client.py | 413 ++++++++--------------- transfer_queue/controller.py | 507 +++++++++++++---------------- transfer_queue/interface.py | 143 ++------ 5 files changed, 436 insertions(+), 688 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c508824a..8b57e5c7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,12 +14,10 @@ and tests without reconstructing hidden control flow. and scheduling; storage managers, clients, and backends move payloads. 4. No over-encapsulation: add an abstraction only for real nontrivial reuse, an important invariant, or an existing sampler/storage interface. -5. Do NOT remove public APIs, backend support, async/distributed behavior, - metadata semantics, defaults, or observability in the name of simplicity. -6. Prioritize behavior reachable through supported APIs, configurations, +5. Prioritize behavior reachable through supported APIs, configurations, backends, tutorials, or tests; do not fix states the system cannot enter. -Details: `.claude/skills/simplicity-first` (invoke before any code change). +Details: `.codex/skills/simplicity-first` (invoke before any code change). ## Comments diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index a1a39761..271758ee 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -212,6 +212,61 @@ def assert_nested_tensor_equal(nested_a, nested_b, msg=""): assert torch.equal(a, b), f"{msg} Component {i} not equal: {a} vs {b}" +class TestRayWorkerKVInterfaceE2E: + @pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) + def test_kv_api_inside_remote_worker(self, tq_system, use_async): + # Pytest test modules are not importable on Ray workers, so keep this local + # and let cloudpickle serialize the function by value. + @ray.remote + def run_kv_api_in_worker(partition_id: str, key: str, use_async: bool) -> dict: + import asyncio + + import torch + + import transfer_queue as tq + + tq.init() + + def call(api_name, *args, **kwargs): + api = getattr(tq, f"async_{api_name}" if use_async else api_name) + result = api(*args, **kwargs) + return asyncio.run(result) if use_async else result + + meta = call( + "kv_put", + key=key, + partition_id=partition_id, + fields={"data": torch.tensor([1, 2, 3])}, + tag={"source": "ray_worker"}, + ) + data = call("kv_batch_get", keys=key, partition_id=partition_id) + keys_before_clear = call("kv_list", partition_id=partition_id) + call("kv_clear", keys=key, partition_id=partition_id) + keys_after_clear = call("kv_list", partition_id=partition_id) + + return { + "fields": meta.fields, + "data": list(data["data"])[0].tolist(), + "listed_before_clear": key in keys_before_clear.get(partition_id, {}), + "listed_after_clear": key in keys_after_clear.get(partition_id, {}), + } + + mode = "async" if use_async else "sync" + result = ray.get( + run_kv_api_in_worker.remote( + partition_id=f"ray_worker_{mode}_partition", + key=f"ray_worker_{mode}_key", + use_async=use_async, + ), + timeout=60, + ) + + assert result["fields"] == ["data"] + assert result["data"] == [1, 2, 3] + assert result["listed_before_clear"] is True + assert result["listed_after_clear"] is False + + class TestKVPutE2E: """End-to-end tests for kv_put functionality.""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index a0b82f92..01e50b85 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -176,6 +176,33 @@ def initialize_storage_manager( zmq_context=self.zmq_context, ) + async def _request_controller( + self, + socket: zmq.asyncio.Socket | None, + request_type: ZMQRequestType, + response_type: ZMQRequestType, + body: dict[str, Any], + ) -> ZMQMessage: + """Send one controller request and validate its response type.""" + assert socket is not None + request_msg = ZMQMessage.create( + request_type=request_type, # type: ignore[arg-type] + sender_id=self.client_id, + receiver_id=self._controller.id, + body=body, + ) + await socket.send_multipart(request_msg.serialize()) + response_serialized = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_serialized) + logger.debug(f"[{self.client_id}]: Received {response_msg.request_type} from controller {self._controller.id}") + if response_msg.request_type != response_type: + message = response_msg.body.get("message", "Unknown error") + raise RuntimeError( + f"[{self.client_id}]: Expected {response_type}, got {response_msg.request_type} " + f"from controller {self._controller.id}: {message}" + ) + return response_msg + # ==================== Basic API ==================== @with_controller_socket async def async_get_meta( @@ -237,11 +264,10 @@ async def async_get_meta( ... )) >>> print(batch_meta.is_ready) # May be False if some samples not ready """ - assert socket is not None - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_META, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_META, + response_type=ZMQRequestType.GET_META_RESPONSE, body={ "data_fields": data_fields, "batch_size": batch_size, @@ -251,21 +277,7 @@ async def async_get_meta( "sampling_config": sampling_config, }, ) - - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get_meta response: {response_msg} from controller {self._controller.id}" - ) - - if response_msg.request_type == ZMQRequestType.GET_META_RESPONSE: - return response_msg.body["metadata"] - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to get metadata from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + return response_msg.body["metadata"] @with_controller_socket async def async_set_custom_meta( @@ -319,28 +331,15 @@ async def async_set_custom_meta( {meta.global_indexes[i]: custom_meta[i] for i in range(len(custom_meta))} ) - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.SET_CUSTOM_META, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.SET_CUSTOM_META, + response_type=ZMQRequestType.SET_CUSTOM_META_RESPONSE, body={ "partition_custom_meta": partition_custom_meta, }, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client set_custom_meta response: {response_msg} from controller {self._controller.id}" - ) - - if response_msg.request_type != ZMQRequestType.SET_CUSTOM_META_RESPONSE: - raise RuntimeError( - f"[{self.client_id}]: Failed to set custom metadata to controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) - async def async_put( self, data: TensorDict, @@ -587,20 +586,13 @@ async def _mark_clearing_in_controller(self, metadata: BatchMeta, socket=None): Raises: RuntimeError: If the controller returns an unexpected response """ - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.MARK_CLEARING, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.MARK_CLEARING, + response_type=ZMQRequestType.MARK_CLEARING_RESPONSE, body={"global_indexes": metadata.global_indexes, "partition_ids": metadata.partition_ids}, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - - if response_msg.request_type != ZMQRequestType.MARK_CLEARING_RESPONSE: - raise RuntimeError("Failed to mark samples as clearing in controller.") - @with_controller_socket async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None): """Clear metadata in the controller. @@ -613,20 +605,13 @@ async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None): RuntimeError: If clear operation fails """ - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_META, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.CLEAR_META, + response_type=ZMQRequestType.CLEAR_META_RESPONSE, body={"global_indexes": metadata.global_indexes, "partition_ids": metadata.partition_ids}, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - - if response_msg.request_type != ZMQRequestType.CLEAR_META_RESPONSE: - raise RuntimeError("Failed to clear samples metadata in controller.") - @with_controller_socket async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta: """Get metadata required for the whole partition from controller. @@ -641,20 +626,12 @@ async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta Raises: RuntimeError: If controller returns error response """ - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_PARTITION_META, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_PARTITION_META, + response_type=ZMQRequestType.GET_PARTITION_META_RESPONSE, body={"partition_id": partition_id}, ) - - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - - if response_msg.request_type != ZMQRequestType.GET_PARTITION_META_RESPONSE: - raise RuntimeError("Failed to get metadata for clear operation.") - return response_msg.body["metadata"] @with_controller_socket @@ -669,20 +646,13 @@ async def _clear_partition_in_controller(self, partition_id, socket=None): RuntimeError: If clear operation fails """ - request_msg = ZMQMessage.create( + await self._request_controller( + socket=socket, request_type=ZMQRequestType.CLEAR_PARTITION, - sender_id=self.client_id, - receiver_id=self._controller.id, + response_type=ZMQRequestType.CLEAR_PARTITION_RESPONSE, body={"partition_id": partition_id}, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - - if response_msg.request_type != ZMQRequestType.CLEAR_PARTITION_RESPONSE: - raise RuntimeError(f"Failed to clear partition {partition_id} in controller.") - # ==================== Status Query API ==================== @with_controller_socket async def async_get_consumption_status( @@ -715,35 +685,19 @@ async def async_get_consumption_status( >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") """ - assert socket is not None - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_CONSUMPTION, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={ - "partition_id": partition_id, - "task_name": task_name, - }, - ) - try: - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get consumption response: {response_msg} " - f"from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_CONSUMPTION, + response_type=ZMQRequestType.CONSUMPTION_RESPONSE, + body={ + "partition_id": partition_id, + "task_name": task_name, + }, ) - - if response_msg.request_type == ZMQRequestType.CONSUMPTION_RESPONSE: - global_index = response_msg.body.get("global_index") - consumption_status = response_msg.body.get("consumption_status") - return global_index, consumption_status - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to get consumption status from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + global_index = response_msg.body.get("global_index") + consumption_status = response_msg.body.get("consumption_status") + return global_index, consumption_status except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in get_consumption_status: {str(e)}") from e @@ -777,35 +731,19 @@ async def async_get_production_status( ... )) >>> print(f"Global index: {global_index}, Production status: {production_status}") """ - assert socket is not None - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_PRODUCTION, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={ - "partition_id": partition_id, - "data_fields": data_fields, - }, - ) - try: - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get production response: {response_msg} " - f"from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_PRODUCTION, + response_type=ZMQRequestType.PRODUCTION_RESPONSE, + body={ + "partition_id": partition_id, + "data_fields": data_fields, + }, ) - - if response_msg.request_type == ZMQRequestType.PRODUCTION_RESPONSE: - global_index = response_msg.body.get("global_index") - production_status = response_msg.body.get("production_status") - return global_index, production_status - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to get production status from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + global_index = response_msg.body.get("global_index") + production_status = response_msg.body.get("production_status") + return global_index, production_status except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in get_data_production_status: {str(e)}") from e @@ -910,34 +848,20 @@ async def async_reset_consumption( ... )) >>> print(f"Reset successful: {success}") """ - assert socket is not None body = {"partition_id": partition_id} if task_name is not None: body["task_name"] = task_name - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.RESET_CONSUMPTION, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body=body, - ) try: - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client reset consumption response: {response_msg} " - f"from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.RESET_CONSUMPTION, + response_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, + body=body, ) - if response_msg.request_type == ZMQRequestType.RESET_CONSUMPTION_RESPONSE: - success = response_msg.body.get("success", False) - if not success: - logger.warning(f"[{self.client_id}]: Reset consumption failed: {response_msg.body.get('message')}") - return success - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to reset consumption from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + success = response_msg.body.get("success", False) + if not success: + logger.warning(f"[{self.client_id}]: Reset consumption failed: {response_msg.body.get('message')}") + return success except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in reset_consumption: {str(e)}") from e @@ -958,31 +882,14 @@ async def async_get_partition_list( >>> partition_ids = asyncio.run(client.get_partition_list()) >>> print(f"Available partitions: {partition_ids}") """ - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_LIST_PARTITIONS, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={}, - ) - try: - assert socket is not None - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get partition list response: {response_msg} " - f"from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.GET_LIST_PARTITIONS, + response_type=ZMQRequestType.LIST_PARTITIONS_RESPONSE, + body={}, ) - - if response_msg.request_type == ZMQRequestType.LIST_PARTITIONS_RESPONSE: - partition_ids = response_msg.body.get("partition_ids", []) - return partition_ids - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to get partition list from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + return response_msg.body.get("partition_ids", []) except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in get_partition_list: {str(e)}") from e @@ -1020,33 +927,18 @@ async def async_kv_retrieve_meta( else: raise TypeError("Only string or list of strings are allowed as `keys`.") - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_META, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={ - "keys": keys, - "partition_id": partition_id, - "create": create, - }, - ) - try: - assert socket is not None, "Socket must be initialized before use" - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}] Received KV_RETRIEVE_META response: {response_msg} " - f"from controller {self._controller.id}" - ) - - if response_msg.request_type == ZMQRequestType.KV_RETRIEVE_META_RESPONSE: - return response_msg.body.get("metadata", BatchMeta.empty()) - - raise RuntimeError( - f"[{self.client_id}] Failed to retrieve metadata {response_msg.body.get('message', 'Unknown error')}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_RETRIEVE_META, + response_type=ZMQRequestType.KV_RETRIEVE_META_RESPONSE, + body={ + "keys": keys, + "partition_id": partition_id, + "create": create, + }, ) + return response_msg.body.get("metadata", BatchMeta.empty()) except Exception as e: raise RuntimeError(f"[{self.client_id}] Failed in async_kv_retrieve_meta: {e}") from e @@ -1083,33 +975,17 @@ async def async_kv_retrieve_keys( else: raise TypeError("Only int or list of int are allowed as `global_indexes`.") - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_KEYS, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={"global_indexes": global_indexes, "partition_id": partition_id}, - ) - try: - assert socket is not None - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get kv_retrieve_indexes response: {response_msg} " - f"from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_RETRIEVE_KEYS, + response_type=ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, + body={"global_indexes": global_indexes, "partition_id": partition_id}, ) - - if response_msg.request_type == ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE: - keys = response_msg.body.get("keys", []) - if len(keys) != len(global_indexes): - raise RuntimeError("Some global_indexes have no corresponding keys!") - return keys - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to retrieve indexes from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + keys = response_msg.body.get("keys", []) + if len(keys) != len(global_indexes): + raise RuntimeError("Some global_indexes have no corresponding keys!") + return keys except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in kv_retrieve_indexes: {str(e)}") from e @@ -1142,32 +1018,14 @@ async def async_kv_list( } """ - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_LIST, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, - body={ - "partition_id": partition_id, - }, - ) - try: - assert socket is not None - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - logger.debug( - f"[{self.client_id}]: Client get kv_list response: {response_msg} from controller {self._controller.id}" + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_LIST, + response_type=ZMQRequestType.KV_LIST_RESPONSE, + body={"partition_id": partition_id}, ) - - if response_msg.request_type == ZMQRequestType.KV_LIST_RESPONSE: - partition_info = response_msg.body.get("partition_info", {}) - return partition_info - else: - raise RuntimeError( - f"[{self.client_id}]: Failed to list keys from controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + return response_msg.body.get("partition_info", {}) except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in kv_list: {str(e)}") from e @@ -1248,21 +1106,12 @@ async def async_save_controller_checkpoint( RuntimeError: If the RPC fails or an unexpected response is received. """ try: - assert socket is not None - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT, + response_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, body={"path": path}, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - if response_msg.request_type != ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE: - raise RuntimeError( - f"[{self.client_id}]: Unexpected response type {response_msg.request_type} " - f"from controller during checkpoint dump" - ) except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in save_controller_checkpoint: {str(e)}") from e @@ -1288,21 +1137,12 @@ async def async_load_controller_checkpoint( RuntimeError: If the RPC fails or an unexpected response is received. """ try: - assert socket is not None - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT, # type: ignore[arg-type] - sender_id=self.client_id, - receiver_id=self._controller.id, + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT, + response_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, body={"path": path}, ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_serialized) - if response_msg.request_type != ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE: - raise RuntimeError( - f"[{self.client_id}]: Unexpected response type {response_msg.request_type} " - f"from controller during checkpoint restore" - ) except Exception as e: raise RuntimeError(f"[{self.client_id}]: Error in load_controller_checkpoint: {str(e)}") from e @@ -1395,18 +1235,19 @@ def _start_loop(self): asyncio.set_event_loop(self._loop) self._loop.run_forever() + def _run_coroutine(self, coro): + """Run a coroutine on the client's background event loop.""" + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + return future.result() + def _bind_sync_methods( self, ): """Convert and bind synchronous methods.""" - def _run(coro): - future = asyncio.run_coroutine_threadsafe(coro, self._loop) - return future.result() - def _make_sync(async_method): def wrapper(*args, **kwargs): - return _run(async_method(*args, **kwargs)) + return self._run_coroutine(async_method(*args, **kwargs)) return wrapper diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 719b4d71..05dea6a1 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1918,302 +1918,253 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | Whatever the matching handler raises propagates to the caller, which turns it into an error response. """ - response_msg = None - - if request_msg.request_type == ZMQRequestType.GET_META: - with monitor.measure(op_type="GET_META"): - params = request_msg.body - - metadata = self.get_metadata( - data_fields=params["data_fields"], - batch_size=params["batch_size"], - partition_id=params["partition_id"], - mode=params.get("mode", "fetch"), - task_name=params.get("task_name"), - sampling_config=params.get("sampling_config", {}), - ) - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) - - elif request_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE: - with monitor.measure(op_type="NOTIFY_DATA_UPDATE"): - message_data = request_msg.body - partition_id = message_data.get("partition_id") - global_indexes = message_data.get("global_indexes", []) - - # Update production status - success = self.update_production_status( - partition_id=cast(str, partition_id), - global_indexes=global_indexes, - field_schema=message_data.get("field_schema", {}), - custom_backend_meta=message_data.get("custom_backend_meta", {}), - ) - if success: - if self._metrics is not None: - self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes)) - logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}") - - # Send acknowledgment - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "controller_id": self.controller_id, - "partition_id": partition_id, - "success": success, - }, - ) - - elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META: - with monitor.measure(op_type="GET_PARTITION_META"): - params = request_msg.body - partition_id = params["partition_id"] - partition = self._get_partition(partition_id) - if partition is not None: - partition_data_fields = list(partition.field_name_mapping.keys()) - - metadata = self.get_metadata( - data_fields=partition_data_fields, - partition_id=partition_id, - mode="force_fetch", - ) - else: - metadata = None - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_PARTITION_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) - elif request_msg.request_type == ZMQRequestType.SET_CUSTOM_META: - with monitor.measure(op_type="SET_CUSTOM_META"): - params = request_msg.body - partition_custom_meta = params["partition_custom_meta"] - - self.set_custom_meta(partition_custom_meta=partition_custom_meta) - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.SET_CUSTOM_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": "Successfully set custom_meta"}, - ) - - elif request_msg.request_type == ZMQRequestType.MARK_CLEARING: - with monitor.measure(op_type="MARK_CLEARING"): - params = request_msg.body - self.mark_clearing(params["global_indexes"], params["partition_ids"]) - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.MARK_CLEARING_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": "Mark clearing completed"}, - ) - - elif request_msg.request_type == ZMQRequestType.CLEAR_META: - with monitor.measure(op_type="CLEAR_META"): - params = request_msg.body - global_indexes = params["global_indexes"] - partition_ids = params["partition_ids"] - - self.clear_meta(global_indexes, partition_ids) - if self._metrics is not None: - self._metrics.record_samples("CLEAR_META", len(global_indexes)) - - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": f"Clear samples operation completed by controller {self.controller_id}"}, - ) - - elif request_msg.request_type == ZMQRequestType.CLEAR_PARTITION: - with monitor.measure(op_type="CLEAR_PARTITION"): - params = request_msg.body - partition_id = params["partition_id"] + handlers = { + ZMQRequestType.GET_META: self._handle_get_meta_request, + ZMQRequestType.NOTIFY_DATA_UPDATE: self._handle_notify_data_update_request, + ZMQRequestType.GET_PARTITION_META: self._handle_get_partition_meta_request, + ZMQRequestType.SET_CUSTOM_META: self._handle_set_custom_meta_request, + ZMQRequestType.MARK_CLEARING: self._handle_mark_clearing_request, + ZMQRequestType.CLEAR_META: self._handle_clear_meta_request, + ZMQRequestType.CLEAR_PARTITION: self._handle_clear_partition_request, + ZMQRequestType.GET_CONSUMPTION: self._handle_get_consumption_request, + ZMQRequestType.RESET_CONSUMPTION: self._handle_reset_consumption_request, + ZMQRequestType.GET_PRODUCTION: self._handle_get_production_request, + ZMQRequestType.GET_LIST_PARTITIONS: self._handle_get_list_partitions_request, + ZMQRequestType.KV_RETRIEVE_META: self._handle_kv_retrieve_meta_request, + ZMQRequestType.KV_RETRIEVE_KEYS: self._handle_kv_retrieve_keys_request, + ZMQRequestType.KV_LIST: self._handle_kv_list_request, + ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, + ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, + } + handler = handlers.get(request_msg.request_type) + if handler is None: + return None - self.clear_partition(partition_id) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_PARTITION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": f"Clear partition operation completed by controller {self.controller_id}"}, - ) + checkpoint_requests = { + ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT, + ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT, + } + if request_msg.request_type in checkpoint_requests: + return handler(request_msg) - elif request_msg.request_type == ZMQRequestType.GET_CONSUMPTION: - with monitor.measure(op_type="GET_CONSUMPTION"): - # Handle consumption status checks - params = request_msg.body + with monitor.measure(op_type=request_msg.request_type.value): + return handler(request_msg) - global_index, consumption_status = self.get_consumption_status( - params["partition_id"], params["task_name"] - ) - sample_filter = params.get("sample_filter") # TODO: DEPRECATED in future + def _make_response( + self, + request_msg: ZMQMessage, + response_type: ZMQRequestType, + body: dict[str, Any], + ) -> ZMQMessage: + """Build a controller response addressed to the request sender.""" + return ZMQMessage.create( + request_type=response_type, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body=body, + ) - if sample_filter and consumption_status is not None: - # TODO: DEPRECATED in future - consumption_status = consumption_status[sample_filter] + def _handle_get_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + metadata = self.get_metadata( + data_fields=params["data_fields"], + batch_size=params["batch_size"], + partition_id=params["partition_id"], + mode=params.get("mode", "fetch"), + task_name=params.get("task_name"), + sampling_config=params.get("sampling_config", {}), + ) + return self._make_response(request_msg, ZMQRequestType.GET_META_RESPONSE, {"metadata": metadata}) + + def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + partition_id = params.get("partition_id") + global_indexes = params.get("global_indexes", []) + success = self.update_production_status( + partition_id=cast(str, partition_id), + global_indexes=global_indexes, + field_schema=params.get("field_schema", {}), + custom_backend_meta=params.get("custom_backend_meta", {}), + ) + if success: + if self._metrics is not None: + self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes)) + logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}") + + return self._make_response( + request_msg, + ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + { + "controller_id": self.controller_id, + "partition_id": partition_id, + "success": success, + }, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": params["partition_id"], - "global_index": global_index, - "consumption_status": consumption_status, - }, - ) + def _handle_get_partition_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: + partition_id = request_msg.body["partition_id"] + partition = self._get_partition(partition_id) + if partition is None: + metadata = None + else: + metadata = self.get_metadata( + data_fields=list(partition.field_name_mapping.keys()), + partition_id=partition_id, + mode="force_fetch", + ) + return self._make_response( + request_msg, + ZMQRequestType.GET_PARTITION_META_RESPONSE, + {"metadata": metadata}, + ) - elif request_msg.request_type == ZMQRequestType.RESET_CONSUMPTION: - with monitor.measure(op_type="RESET_CONSUMPTION"): - # Handle reset consumption status request - params = request_msg.body - partition_id = params["partition_id"] - task_name = params.get("task_name") # Optional - try: - self.reset_consumption(partition_id, task_name) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": partition_id, - "success": True, - "message": f"Consumption reset for partition {partition_id}", - }, - ) - except Exception as e: - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": partition_id, - "success": False, - "message": str(e), - }, - ) + def _handle_set_custom_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: + self.set_custom_meta(partition_custom_meta=request_msg.body["partition_custom_meta"]) + return self._make_response( + request_msg, + ZMQRequestType.SET_CUSTOM_META_RESPONSE, + {"message": "Successfully set custom_meta"}, + ) - elif request_msg.request_type == ZMQRequestType.GET_PRODUCTION: - with monitor.measure(op_type="GET_PRODUCTION"): - # Handle production status checks - params = request_msg.body + def _handle_mark_clearing_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + self.mark_clearing(params["global_indexes"], params["partition_ids"]) + return self._make_response( + request_msg, + ZMQRequestType.MARK_CLEARING_RESPONSE, + {"message": "Mark clearing completed"}, + ) - global_index, production_status = self.get_production_status( - params["partition_id"], params["data_fields"] - ) + def _handle_clear_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + global_indexes = params["global_indexes"] + self.clear_meta(global_indexes, params["partition_ids"]) + if self._metrics is not None: + self._metrics.record_samples("CLEAR_META", len(global_indexes)) + return self._make_response( + request_msg, + ZMQRequestType.CLEAR_META_RESPONSE, + {"message": f"Clear samples operation completed by controller {self.controller_id}"}, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PRODUCTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": params["partition_id"], - "global_index": global_index, - "production_status": production_status, - }, - ) + def _handle_clear_partition_request(self, request_msg: ZMQMessage) -> ZMQMessage: + partition_id = request_msg.body["partition_id"] + self.clear_partition(partition_id) + return self._make_response( + request_msg, + ZMQRequestType.CLEAR_PARTITION_RESPONSE, + {"message": f"Clear partition operation completed by controller {self.controller_id}"}, + ) - elif request_msg.request_type == ZMQRequestType.GET_LIST_PARTITIONS: - with monitor.measure(op_type="GET_LIST_PARTITIONS"): - # Handle list partitions request - partition_ids = self.list_partitions() - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.LIST_PARTITIONS_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"partition_ids": partition_ids}, - ) + def _handle_get_consumption_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + global_index, consumption_status = self.get_consumption_status(params["partition_id"], params["task_name"]) + sample_filter = params.get("sample_filter") # TODO: DEPRECATED in future + if sample_filter and consumption_status is not None: + consumption_status = consumption_status[sample_filter] # TODO: DEPRECATED in future + + return self._make_response( + request_msg, + ZMQRequestType.CONSUMPTION_RESPONSE, + { + "partition_id": params["partition_id"], + "global_index": global_index, + "consumption_status": consumption_status, + }, + ) - elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_META: - with monitor.measure(op_type="KV_RETRIEVE_META"): - params = request_msg.body - keys = params["keys"] - partition_id = params["partition_id"] - create = params["create"] - - metadata = self.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) + def _handle_reset_consumption_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + partition_id = params["partition_id"] + try: + self.reset_consumption(partition_id, params.get("task_name")) + success = True + message = f"Consumption reset for partition {partition_id}" + except Exception as e: + success = False + message = str(e) + return self._make_response( + request_msg, + ZMQRequestType.RESET_CONSUMPTION_RESPONSE, + {"partition_id": partition_id, "success": success, "message": message}, + ) - elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_KEYS: - with monitor.measure(op_type="KV_RETRIEVE_KEYS"): - params = request_msg.body - global_indexes = params["global_indexes"] - partition_id = params["partition_id"] + def _handle_get_production_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + global_index, production_status = self.get_production_status(params["partition_id"], params["data_fields"]) + return self._make_response( + request_msg, + ZMQRequestType.PRODUCTION_RESPONSE, + { + "partition_id": params["partition_id"], + "global_index": global_index, + "production_status": production_status, + }, + ) - keys = self.kv_retrieve_keys(global_indexes=global_indexes, partition_id=partition_id) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"keys": keys}, - ) + def _handle_get_list_partitions_request(self, request_msg: ZMQMessage) -> ZMQMessage: + return self._make_response( + request_msg, + ZMQRequestType.LIST_PARTITIONS_RESPONSE, + {"partition_ids": self.list_partitions()}, + ) - elif request_msg.request_type == ZMQRequestType.KV_LIST: - with monitor.measure(op_type="KV_LIST"): - params = request_msg.body - partition_id = params["partition_id"] - if partition_id is None: - partition_id = list(self.partitions.keys()) - else: - partition_id = [partition_id] - - message = "success" - partition_info = {} - for pid in partition_id: - partition = self._get_partition(pid) - if partition: - keys = list(partition.keys_mapping.keys()) - single_partition_info = { - k: partition.custom_meta.get(partition.keys_mapping[k], {}) for k in keys - } - partition_info[pid] = single_partition_info - else: - # this only happens when params["partition_id"] is not None - message = f"partition {pid} does not exist" + def _handle_kv_retrieve_meta_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + metadata = self.kv_retrieve_meta( + keys=params["keys"], + partition_id=params["partition_id"], + create=params["create"], + ) + return self._make_response( + request_msg, + ZMQRequestType.KV_RETRIEVE_META_RESPONSE, + {"metadata": metadata}, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_LIST_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"partition_info": partition_info, "message": message}, - ) + def _handle_kv_retrieve_keys_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + keys = self.kv_retrieve_keys( + global_indexes=params["global_indexes"], + partition_id=params["partition_id"], + ) + return self._make_response(request_msg, ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, {"keys": keys}) + + def _handle_kv_list_request(self, request_msg: ZMQMessage) -> ZMQMessage: + requested_partition_id = request_msg.body["partition_id"] + partition_ids = list(self.partitions.keys()) if requested_partition_id is None else [requested_partition_id] + message = "success" + partition_info = {} + for partition_id in partition_ids: + partition = self._get_partition(partition_id) + if partition is None: + message = f"partition {partition_id} does not exist" + continue + partition_info[partition_id] = { + key: partition.custom_meta.get(index, {}) for key, index in partition.keys_mapping.items() + } - elif request_msg.request_type == ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: - path = request_msg.body["path"] - self.save_checkpoint(path) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"success": True}, - ) + return self._make_response( + request_msg, + ZMQRequestType.KV_LIST_RESPONSE, + {"partition_info": partition_info, "message": message}, + ) - elif request_msg.request_type == ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: - path = request_msg.body["path"] - self.load_checkpoint(path) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"success": True}, - ) + def _handle_save_controller_checkpoint_request(self, request_msg: ZMQMessage) -> ZMQMessage: + self.save_checkpoint(request_msg.body["path"]) + return self._make_response( + request_msg, + ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, + {"success": True}, + ) - return response_msg + def _handle_load_controller_checkpoint_request(self, request_msg: ZMQMessage) -> ZMQMessage: + self.load_checkpoint(request_msg.body["path"]) + return self._make_response( + request_msg, + ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, + {"success": True}, + ) def get_zmq_server_info(self) -> ZMQServerInfo: """Get ZMQ server connection information.""" diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index fd129560..cce2abdd 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -371,51 +371,15 @@ def kv_put( ... ) >>> print(meta.fields) # ['input_ids'] """ - if fields is None and tag is None: - raise ValueError("Please provide at least one parameter of `fields` or `tag`.") - tq_client = _maybe_create_tq_client() - - # 1. translate user-specified key to BatchMeta - batch_meta = tq_client.kv_retrieve_meta(keys=[key], partition_id=partition_id, create=True) - - if batch_meta.size != 1: - raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match with input `key` size of 1!") - - # 2. register the user-specified tag to BatchMeta - if tag is not None: - batch_meta.update_custom_meta([tag]) - - # 3. put data - if fields is not None: - if isinstance(fields, dict): - # TODO: consider whether to support this... - batch = {} - for field_name, value in fields.items(): - if isinstance(value, torch.Tensor): - if value.is_nested: - raise ValueError("Please use (async)kv_batch_put for batch operation") - batch[field_name] = value.unsqueeze(0) - else: - batch[field_name] = NonTensorStack(value) - fields = TensorDict(batch, batch_size=[1]) - elif not isinstance(fields, TensorDict): - raise ValueError("`fields` can only be dict or TensorDict") - - # After put, batch_meta.field_names will include the new fields written by user - batch_meta = tq_client.put(fields, batch_meta, data_parser=data_parser) - else: - # Directly update custom_meta (tag) to controller - tq_client.set_custom_meta(batch_meta) - - fields_to_return = batch_meta.field_names - - return KVBatchMeta( - keys=[key], - tags=batch_meta.custom_meta, - partition_id=partition_id, - fields=fields_to_return, - extra_info=batch_meta.extra_info, + return tq_client._run_coroutine( + async_kv_put( + key=key, + partition_id=partition_id, + fields=fields, + tag=tag, + data_parser=data_parser, + ) ) @@ -466,36 +430,15 @@ def kv_batch_put( >>> meta = tq.kv_batch_put(keys=keys, partition_id="train", fields=fields, tags=tags) >>> print(meta.fields) """ - num_keys = len(keys) - - if fields is None and tags is None: - raise ValueError("Please provide at least one parameter of fields or tag.") - - if fields is not None and fields.batch_size[0] != num_keys: - raise ValueError(f"Length of `keys` ({num_keys}) does not match `fields` batch size ({fields.batch_size[0]}).") - tq_client = _maybe_create_tq_client() - batch_meta = tq_client.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True) - - if batch_meta.size != num_keys: - raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match input `keys` size {num_keys}.") - - if tags is not None: - if len(tags) != num_keys: - raise ValueError(f"Length of `keys` ({num_keys}) does not match length of `tags` ({len(tags)}).") - batch_meta.update_custom_meta(tags) - - if fields is not None: - batch_meta = tq_client.put(fields, batch_meta, data_parser=data_parser) - else: # tags is not None - tq_client.set_custom_meta(batch_meta) - - return KVBatchMeta( - keys=keys, - tags=batch_meta.custom_meta, - partition_id=partition_id, - fields=batch_meta.field_names, - extra_info=batch_meta.extra_info, + return tq_client._run_coroutine( + async_kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=fields, + tags=tags, + data_parser=data_parser, + ) ) @@ -535,23 +478,8 @@ def kv_batch_get_by_meta(meta: KVBatchMeta, select_fields: list[str] | str | Non >>> # Then retrieve it using the returned metadata >>> data = tq.kv_batch_get_by_meta(meta) """ - if meta.partition_id is None: - raise ValueError("Must provide partition_id in the input KVBatchMeta.") - if select_fields is not None: - if isinstance(select_fields, str): - fields_to_fetch: list[str] | None = [select_fields] - else: - fields_to_fetch = select_fields - - assert fields_to_fetch is not None - if meta.fields is None or any(f not in meta.fields for f in fields_to_fetch): - raise ValueError( - f"Some fields assigned in select_fields not found in the metadata. " - f"Assigned: {fields_to_fetch}; Fields in KVBatchMeta: {meta.fields}." - ) - else: - fields_to_fetch = meta.fields - return kv_batch_get(keys=meta.keys, partition_id=meta.partition_id, select_fields=fields_to_fetch) + tq_client = _maybe_create_tq_client() + 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: @@ -584,25 +512,9 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s ... ) """ tq_client = _maybe_create_tq_client() - - batch_meta = tq_client.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) - - if batch_meta.size == 0: - raise ValueError("keys or partition were not found!") - - fields_to_fetch: list[str] | None - if select_fields is not None: - if isinstance(select_fields, str): - fields_to_fetch = [select_fields] - else: - fields_to_fetch = select_fields - batch_meta = batch_meta.select_fields(fields_to_fetch) - - if not batch_meta.is_ready: - raise ValueError("Some fields are not ready in all the requested keys!") - - data = tq_client.get_data(batch_meta) - return data + return tq_client._run_coroutine( + async_kv_batch_get(keys=keys, partition_id=partition_id, select_fields=select_fields) + ) def kv_list(partition_id: str | None = None) -> dict[str, dict[str, Any]]: @@ -640,10 +552,7 @@ def kv_list(partition_id: str | None = None) -> dict[str, dict[str, Any]]: >>> print(f"Partition: {pid}, Key count: {len(keys)}") """ tq_client = _maybe_create_tq_client() - - partition_info = tq_client.kv_list(partition_id) - - return partition_info + return tq_client._run_coroutine(async_kv_list(partition_id=partition_id)) def kv_clear(keys: list[str] | str, partition_id: str) -> None: @@ -665,14 +574,8 @@ def kv_clear(keys: list[str] | str, partition_id: str) -> None: >>> tq.kv_clear(keys=["sample_1", "sample_2"], partition_id="train") """ - if isinstance(keys, str): - keys = [keys] - tq_client = _maybe_create_tq_client() - batch_meta = tq_client.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) - - if batch_meta.size > 0: - tq_client.clear_samples(batch_meta) + tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id)) # ==================== KV Interface API ====================