diff --git a/recipe/simple_use_case/async_demo.py b/recipe/simple_use_case/async_demo.py index 8727cf8..507ee01 100644 --- a/recipe/simple_use_case/async_demo.py +++ b/recipe/simple_use_case/async_demo.py @@ -306,7 +306,7 @@ def fit(self): # Client notifies controller to clear data status, controller returns metadata; # Client then notifies the storage plane to clear based on metadata - asyncio.run(self.data_system_client.async_clear(partition_id=f"train_{step}")) + asyncio.run(self.data_system_client.async_clear_partition(partition_id=f"train_{step}")) logger.info("clear ok! ") logger.info("demo done!") diff --git a/recipe/simple_use_case/sync_demo.py b/recipe/simple_use_case/sync_demo.py index 1c03031..f559df0 100644 --- a/recipe/simple_use_case/sync_demo.py +++ b/recipe/simple_use_case/sync_demo.py @@ -190,7 +190,7 @@ def fit(config, data_system_client): # Client then notifies the storage plane to clear based on metadata # Client selects one master controller to get metadata, # other controllers directly clear without returning metadata - data_system_client.clear(partition_id=f"train_{step}") + data_system_client.clear_partition(partition_id=f"train_{step}") logger.info("clear ok! ") logger.info("demo done!") diff --git a/tests/test_client.py b/tests/test_client.py index 93a3696..7632d42 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -101,12 +101,16 @@ def _handle_requests(self): if request_msg.request_type == ZMQRequestType.GET_META: response_body = self._mock_batch_meta(request_msg.body) response_type = ZMQRequestType.GET_META_RESPONSE - elif request_msg.request_type == ZMQRequestType.GET_CLEAR_META: - response_body = self._mock_batch_meta(request_msg.body) - response_type = ZMQRequestType.GET_CLEAR_META_RESPONSE elif request_msg.request_type == ZMQRequestType.CLEAR_META: - response_body = {"message": "clear ok"} + response_body = {"message": "clear meta ok"} response_type = ZMQRequestType.CLEAR_META_RESPONSE + elif request_msg.request_type == ZMQRequestType.CLEAR_PARTITION: + response_body = {"message": "clear partition ok"} + response_type = ZMQRequestType.CLEAR_PARTITION_RESPONSE + elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META: + # Mock partition metadata response + response_body = {"metadata": self._mock_batch_meta(request_msg.body)} + response_type = ZMQRequestType.GET_PARTITION_META_RESPONSE elif request_msg.request_type == ZMQRequestType.CHECK_CONSUMPTION: # Mock consumption status check - all consumed response_body = { @@ -375,14 +379,6 @@ def test_get_meta(client_setup): assert len(metadata.global_indexes) == 10 -def test_clear_operation(client_setup): - """Test clear operation""" - client, _, _ = client_setup - - # Test clear operation - client.clear(partition_id="0") - - # Test with single controller and multiple storage units def test_single_controller_multiple_storages(): """Test client with single controller and multiple storage units""" @@ -517,3 +513,72 @@ async def test_async_get_partition_list(client_setup): assert "partition_0" in partition_list assert "partition_1" in partition_list assert "test_partition" in partition_list + + +# Test clear methods +@pytest.mark.asyncio +async def test_async_clear_partition(client_setup): + """Test async clear partition operation""" + client, _, _ = client_setup + + # Test async_clear_partition + await client.async_clear_partition(partition_id="test_partition") + + # If no exception is raised, the test passes + assert True + + +@pytest.mark.asyncio +async def test_async_clear_samples(client_setup): + """Test async clear samples operation""" + client, _, _ = client_setup + + # First get metadata to create a BatchMeta object + metadata = await client.async_get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0") + + # Test async_clear_samples + await client.async_clear_samples(metadata=metadata) + + # If no exception is raised, the test passes + assert True + + +def test_clear_partition(client_setup): + """Test synchronous clear partition operation""" + client, _, _ = client_setup + + # Test synchronous clear_partition + client.clear_partition(partition_id="test_partition") + + # If no exception is raised, the test passes + assert True + + +def test_clear_samples(client_setup): + """Test synchronous clear samples operation""" + client, _, _ = client_setup + + # First get metadata to create a BatchMeta object + metadata = client.get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0") + + # Test synchronous clear_samples + client.clear_samples(metadata=metadata) + + # If no exception is raised, the test passes + assert True + + +@pytest.mark.asyncio +async def test_async_clear_samples_with_empty_metadata(client_setup): + """Test async_clear_samples with empty BatchMeta""" + client, _, _ = client_setup + + # Create empty BatchMeta + metadata = BatchMeta(samples=[]) + + # The clear operation should complete without raising an exception + # because the mock storage manager is configured to handle this + await client.async_clear_samples(metadata=metadata) + + # If no exception is raised, the test passes + assert True diff --git a/tests/test_controller.py b/tests/test_controller.py index 9585d99..a036f2c 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -173,13 +173,13 @@ def test_controller_with_single_partition(self, ray_setup): assert [sample.fields for sample in clear_meta.samples] == [{}] * (gbs * num_n_samples) print("✓ Clear metadata correct") - # Test clear - ray.get(tq_controller.clear.remote(partition_id)) + # Test clear_partition + ray.get(tq_controller.clear_partition.remote(partition_id)) partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id)) partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id)) assert partition_index_range == set() assert partition is None - print("✓ Clear correct") + print("✓ Clear partition correct") def test_controller_with_multi_partitions(self, ray_setup): gbs_1 = 8 @@ -284,7 +284,7 @@ def test_controller_with_multi_partitions(self, ray_setup): # Clear partition 1 partition_index_range_1 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1)) assert partition_index_range_1 - ray.get(tq_controller.clear.remote(partition_id_1)) + ray.get(tq_controller.clear_partition.remote(partition_id_1)) partition_1_after_clear = ray.get(tq_controller.get_partition_snapshot.remote(partition_id_1)) partition_index_range_1_after_clear = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1)) @@ -320,3 +320,64 @@ def test_controller_with_multi_partitions(self, ray_setup): partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id_3)) assert partition_index_range == set(list(range(32)) + list(range(48, 80))) print("✓ Correctly assign partition_3") + + def test_controller_clear_meta(self, ray_setup): + """Test clear_meta functionality for individual samples""" + gbs = 4 + num_n_samples = 2 + partition_id = "test_clear_meta" + + tq_controller = TransferQueueController.remote() + + # Create metadata in insert mode + data_fields = ["prompt_ids", "attention_mask"] + metadata = ray.get( + tq_controller.get_metadata.remote( + data_fields=data_fields, + batch_size=gbs * num_n_samples, + partition_id=partition_id, + mode="insert", + ) + ) + + assert metadata.global_indexes == list(range(gbs * num_n_samples)) + + # Update production status + dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in metadata.global_indexes} + shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in metadata.global_indexes} + success = ray.get( + tq_controller.update_production_status.remote( + partition_id=partition_id, + global_indexes=metadata.global_indexes, + field_names=metadata.field_names, + dtypes=dtypes, + shapes=shapes, + ) + ) + assert success + + # Get partition snapshot before clear + partition_before = ray.get(tq_controller.get_partition_snapshot.remote(partition_id)) + assert partition_before is not None + assert len(partition_before.global_indexes) == gbs * num_n_samples + assert set(partition_before.global_indexes) == set(range(gbs * num_n_samples)) + + # Test clear_meta - clear first 4 samples (indexes 0-3) + global_indexes_to_clear = [0, 1, 2, 3, 6] + partition_ids_to_clear = [partition_id] * len(global_indexes_to_clear) + + ray.get( + tq_controller.clear_meta.remote( + global_indexes=global_indexes_to_clear, + partition_ids=partition_ids_to_clear, + ) + ) + + # Check that only the cleared samples are affected + partition_after = ray.get(tq_controller.get_partition_snapshot.remote(partition_id)) + assert partition_after is not None + + # Verify production status is cleared for the specified indexes + assert set(partition_after.global_indexes) == set([4, 5, 7]) + + print("✓ Clear meta correct") diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index ba2ae1e..85ffc1e 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -314,8 +314,7 @@ def test_data_partition_status_advanced(): initial_consumption_sum = sum(t.sum().item() for t in partition.consumption_status.values()) # Clear only production data - success = partition.clear_data(list(range(4)), clear_consumption=False) - assert success + partition.clear_data(list(range(4)), clear_consumption=False) assert partition.production_status[:4, :].sum().item() == 0 # Consumption data should remain diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 78fa177..b40e0b7 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -50,7 +50,7 @@ handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s")) logger.addHandler(handler) -TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 16)) +TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) class AsyncTransferQueueClient: @@ -388,8 +388,8 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: return results - async def async_clear(self, partition_id: str): - """Asynchronously clear data from all storage units and controller metadata. + async def async_clear_partition(self, partition_id: str): + """Asynchronously clear the whole partition from all storage units and the controller. Args: partition_id: The partition id to clear data for @@ -407,24 +407,83 @@ async def async_clear(self, partition_id: str): if not self._controller: raise RuntimeError("No controller registered") - metadata = await self._get_clear_meta(partition_id) + metadata = await self._get_partition_meta(partition_id) # Clear the controller metadata - await self._clear_controller(partition_id) + await self._clear_partition_in_controller(partition_id) # Clear storage unit data await self.storage_manager.clear_data(metadata) - logger.info(f"[{self.client_id}]: Clear operation for partition_id {partition_id} completed.") + logger.debug(f"[{self.client_id}]: Clear operation for partition_id {partition_id} completed.") except Exception as e: raise RuntimeError(f"Error in clear operation: {str(e)}") from e + async def async_clear_samples(self, metadata: BatchMeta): + """Asynchronously clear specific samples from all storage units and the controller. + + Args: + metadata: The BatchMeta of the corresponding data to be cleared + + Raises: + RuntimeError: If clear operation fails + """ + try: + if not hasattr(self, "storage_manager") or self.storage_manager is None: + raise RuntimeError( + f"[{self.client_id}]: Storage manager not initialized. " + "Call initialize_storage_manager() before performing storage operations." + ) + + if metadata.size == 0: + logger.warning(f"[{self.client_id}]: Empty BatchMeta provided to clear_samples. No action taken.") + return + + if not self._controller: + raise RuntimeError("No controller registered") + + # Clear the controller metadata + await self._clear_meta_in_controller(metadata) + + # Clear storage unit data + await self.storage_manager.clear_data(metadata) + + logger.debug(f"[{self.client_id}]: Clear operation for batch {metadata} completed.") + except Exception as e: + raise RuntimeError(f"Error in clear_samples operation: {str(e)}") from e + @dynamic_socket(socket_name="request_handle_socket") - async def _get_clear_meta(self, partition_id: str, socket=None) -> BatchMeta: - """Get metadata required for clear operation from controller. + async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None): + """Clear metadata in the controller. Args: - partition_id: Partition id to get clear metadata for + metadata: The BatchMeta of the corresponding data to be cleared + socket: ZMQ socket (injected by decorator) + + Raises: + RuntimeError: If clear operation fails + """ + + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.CLEAR_META, + sender_id=self.client_id, + receiver_id=self._controller.id, + body={"global_indexes": metadata.global_indexes, "partition_ids": metadata.partition_ids}, + ) + + await socket.send_multipart(request_msg.serialize()) + response_serialized = await socket.recv_multipart() + response_msg = ZMQMessage.deserialize(response_serialized) + + if response_msg.request_type != ZMQRequestType.CLEAR_META_RESPONSE: + raise RuntimeError("Failed to clear samples metadata in controller.") + + @dynamic_socket(socket_name="request_handle_socket") + async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta: + """Get metadata required for the whole partition from controller. + + Args: + partition_id: Partition id to get partition metadata for socket: ZMQ socket (injected by decorator) Returns: @@ -434,7 +493,7 @@ async def _get_clear_meta(self, partition_id: str, socket=None) -> BatchMeta: RuntimeError: If controller returns error response """ request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_CLEAR_META, + request_type=ZMQRequestType.GET_PARTITION_META, sender_id=self.client_id, receiver_id=self._controller.id, body={"partition_id": partition_id}, @@ -444,16 +503,14 @@ async def _get_clear_meta(self, partition_id: str, socket=None) -> BatchMeta: response_serialized = await socket.recv_multipart() response_msg = ZMQMessage.deserialize(response_serialized) - if response_msg.request_type != ZMQRequestType.GET_CLEAR_META_RESPONSE: - raise RuntimeError( - f"Failed to get metadata for clear operation: {response_msg.body.get('message', 'Unknown error')}" - ) + if response_msg.request_type != ZMQRequestType.GET_PARTITION_META_RESPONSE: + raise RuntimeError("Failed to get metadata for clear operation.") return response_msg.body["metadata"] @dynamic_socket(socket_name="request_handle_socket") - async def _clear_controller(self, partition_id, socket=None): - """Clear metadata from controller. + async def _clear_partition_in_controller(self, partition_id, socket=None): + """Clear the whole partition in the controller. Args: partition_id: Partition id to clear metadata for @@ -462,31 +519,20 @@ async def _clear_controller(self, partition_id, socket=None): Raises: RuntimeError: If clear operation fails """ - try: - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_META, - sender_id=self.client_id, - receiver_id=self._controller.id, - body={"partition_id": partition_id}, - ) - await socket.send_multipart(request_msg.serialize()) - response_serialized = await socket.recv_multipart() - response_msg = ZMQMessage.deserialize(response_serialized) + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.CLEAR_PARTITION, + sender_id=self.client_id, + receiver_id=self._controller.id, + body={"partition_id": partition_id}, + ) - if response_msg.request_type != ZMQRequestType.CLEAR_META_RESPONSE: - raise RuntimeError( - f"Failed to clear controller {self._controller.id}: " - f"{response_msg.body.get('message', 'Unknown error')}" - ) + await socket.send_multipart(request_msg.serialize()) + response_serialized = await socket.recv_multipart() + response_msg = ZMQMessage.deserialize(response_serialized) - logger.info( - f"[{self.client_id}]: Successfully clear controller {self._controller.id} for partition_id " - f"{partition_id}" - ) - except Exception as e: - logger.error(f"[{self.client_id}]: Error clearing controller {self._controller.id}: {str(e)}") - raise + if response_msg.request_type != ZMQRequestType.CLEAR_PARTITION_RESPONSE: + raise RuntimeError(f"Failed to clear partition {partition_id} in controller.") @dynamic_socket(socket_name="request_handle_socket") async def async_check_consumption_status( @@ -736,13 +782,21 @@ def get_data(self, metadata: BatchMeta) -> TensorDict: """ return asyncio.run(self.async_get_data(metadata)) - def clear(self, partition_id: str): - """Synchronously clear data from storage units and controller metadata. + def clear_partition(self, partition_id: str): + """Synchronously clear the whole partition from storage units and controller. Args: partition_id: The partition id to clear data for """ - return asyncio.run(self.async_clear(partition_id)) + return asyncio.run(self.async_clear_partition(partition_id)) + + def clear_samples(self, metadata: BatchMeta): + """Synchronously clear specific samples from storage units and controller metadata. + + Args: + metadata: The BatchMeta of the corresponding data to be cleared + """ + return asyncio.run(self.async_clear_samples(metadata)) def check_consumption_status(self, task_name: str, partition_id: str) -> bool: """Synchronously check if all samples for a partition have been consumed by a specific task. diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 2ce4db3..e0c9bd6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -18,6 +18,7 @@ import time from collections import defaultdict from dataclasses import dataclass, field +from itertools import groupby from operator import itemgetter from threading import Lock, Thread from typing import Any, Optional @@ -138,7 +139,7 @@ def allocate_indexes(self, partition_id, count=1) -> list: return indexes - def release_indexes(self, partition_id) -> list[int]: + def release_partition(self, partition_id) -> list[int]: """ Release all global_indexes of the specified partition, adding them to reusable pool. @@ -158,9 +159,33 @@ def release_indexes(self, partition_id) -> list[int]: for idx in indexes: self.allocated_indexes.discard(idx) - return indexes + return list(indexes) return [] + def release_indexes(self, partition_id: str, indexes_to_release: list[int]): + """ + Release specific global_indexes for a partition, adding them to reusable pool. + + Args: + partition_id: Partition ID + indexes_to_release: List of specific indexes to release + """ + if partition_id not in self.partition_to_indexes: + return [] + + partition_indexes = self.partition_to_indexes[partition_id] + + if not set(indexes_to_release).issubset(partition_indexes): + raise ValueError("Some indexes to release do not belong to the specified partition.") + + partition_indexes.difference_update(indexes_to_release) + self.reusable_indexes.extend(indexes_to_release) + self.allocated_indexes.difference_update(indexes_to_release) + + # If partition has no more indexes, remove it from the mapping + if not partition_indexes: + self.partition_to_indexes.pop(partition_id, None) + def get_indexes_for_partition(self, partition_id) -> set[int]: """ Get all global_indexes for the specified partition. @@ -595,20 +620,23 @@ def _perform_copy(): else: return _perform_copy() - def clear_data(self, global_indexes_range: list[int], clear_consumption: bool = True) -> bool: - """Clear all production and optionally consumption data.""" + def clear_data(self, indexes_to_release: list[int], clear_consumption: bool = True): + """Clear all production and optionally consumption data for given global_indexes.""" try: if self.production_status is not None: - self.production_status[global_indexes_range, :] = 0 + self.production_status[indexes_to_release, :] = 0 if clear_consumption: for consumption_tensor in self.consumption_status.values(): - consumption_tensor[global_indexes_range] = 0 + consumption_tensor[indexes_to_release] = 0 + + self.global_indexes.difference_update(indexes_to_release) - return True except Exception as e: - logger.error(f"Error clearing data for partition {self.partition_id}: {e}") - return False + logger.error( + f"Error clearing data for partition {self.partition_id}: {e}. " + f"Attempted to clear global_indexes: {indexes_to_release}" + ) @ray.remote(num_cpus=1) @@ -1061,28 +1089,69 @@ def generate_batch_meta( return BatchMeta(samples=samples) - def clear(self, partition_id: str, clear_consumption: bool = True) -> bool: + def clear_partition(self, partition_id: str, clear_consumption: bool = True): """ - Clear data for a specific partition. + Clear data for a specific partition (delete the whole partition). Args: partition_id: ID of the partition to clear clear_consumption: Whether to also clear consumption status - - Returns: - True if cleared successfully, False otherwise """ + + logger.debug(f"Clearing data for partition {partition_id}") + partition = self._get_partition(partition_id) if not partition: raise ValueError(f"Partition {partition_id} not found") global_indexes_range = list(self.index_manager.get_indexes_for_partition(partition_id)) - success = partition.clear_data(global_indexes_range, clear_consumption) - self.index_manager.release_indexes(partition_id) + partition.clear_data(global_indexes_range, clear_consumption) + self.index_manager.release_partition(partition_id) self.partitions.pop(partition_id) - if success: - logger.info(f"Cleared data for partition {partition_id}") - return success + + def clear_meta(self, global_indexes: list[int], partition_ids: list[str], clear_consumption: bool = True): + """ + Clear meta for individual samples (preserving the partition). + + Args: + global_indexes: global_indexes to clear + partition_ids: IDs of the partitions to clear + clear_consumption: Whether to also clear consumption status + """ + + logger.debug( + f"{self.controller_id}: Clearing meta with global_indexes {global_indexes} in partition {partition_ids}" + ) + + if global_indexes is None or partition_ids is None: + raise ValueError("global_indexes and partition_ids cannot be None") + + if len(global_indexes) != len(partition_ids): + raise ValueError( + f"global_indexes and partition_ids must have the same length, " + f"got {len(global_indexes)} and {len(partition_ids)}" + ) + + combined = list(zip(partition_ids, global_indexes, strict=True)) + combined.sort(key=itemgetter(0)) + + for partition_id, group in groupby(combined, key=itemgetter(0)): + partition = self._get_partition(partition_id) + if not partition: + raise ValueError(f"Partition {partition_id} not found") + + global_indexes_to_clear = [idx for _, idx in group] + if not set(global_indexes_to_clear).issubset(partition.global_indexes): + raise ValueError( + f"Some global_indexes to clear do not exist in partition {partition_id}. " + f"Target: {global_indexes_to_clear}, Existing: {partition.global_indexes}" + ) + + # Clear data from partition + partition.clear_data(global_indexes_to_clear, clear_consumption) + + # Release the specific indexes from index manager + self.index_manager.release_indexes(partition_id, global_indexes_to_clear) def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" @@ -1126,7 +1195,7 @@ def _wait_connection(self): poller = zmq.Poller() poller.register(self.handshake_socket, zmq.POLLIN) - logger.info(f"Controller {self.controller_id} started waiting for storage connections...") + logger.debug(f"Controller {self.controller_id} started waiting for storage connections...") while True: socks = dict(poller.poll(TQ_CONTROLLER_CONNECTION_CHECK_INTERVAL * 1000)) @@ -1153,7 +1222,7 @@ def _wait_connection(self): if storage_manager_id not in self._connected_storage_managers: self._connected_storage_managers.add(storage_manager_id) storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown") - logger.info( + logger.debug( f"Controller {self.controller_id} received handshake from " f"storage manager {storage_manager_id} (type: {storage_manager_type}). " f"Total connected: {len(self._connected_storage_managers)}" @@ -1223,8 +1292,8 @@ def _process_request(self): body={"metadata": metadata}, ) - elif request_msg.request_type == ZMQRequestType.GET_CLEAR_META: - with perf_monitor.measure(op_type="GET_CLEAR_META"): + elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META: + with perf_monitor.measure(op_type="GET_PARTITION_META"): params = request_msg.body partition_id = params["partition_id"] @@ -1234,31 +1303,38 @@ def _process_request(self): mode="insert", ) response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_CLEAR_META_RESPONSE, + 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.CLEAR_META: with perf_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) + + 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 perf_monitor.measure(op_type="CLEAR_PARTITION"): params = request_msg.body partition_id = params["partition_id"] - clear_success = self.clear(partition_id) - if clear_success: - 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 operation completed by controller {self.controller_id}"}, - ) - else: - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"error": f"Clear operation failed for partition {partition_id}"}, - ) + 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}"}, + ) elif request_msg.request_type == ZMQRequestType.CHECK_CONSUMPTION: with perf_monitor.measure(op_type="CHECK_CONSUMPTION"): diff --git a/transfer_queue/storage/simple_backend.py b/transfer_queue/storage/simple_backend.py index 2051afd..9dc75ed 100644 --- a/transfer_queue/storage/simple_backend.py +++ b/transfer_queue/storage/simple_backend.py @@ -42,7 +42,7 @@ logger.addHandler(handler) TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5)) # in seconds -TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 16)) +TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) class StorageUnitData: diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c0c2e86..341d294 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -64,10 +64,12 @@ class ZMQRequestType(ExplicitEnum): # META_OPERATION GET_META = "GET_META" GET_META_RESPONSE = "GET_META_RESPONSE" - GET_CLEAR_META = "GET_CLEAR_META" - GET_CLEAR_META_RESPONSE = "GET_CLEAR_META_RESPONSE" + GET_PARTITION_META = "GET_PARTITION_META" + GET_PARTITION_META_RESPONSE = "GET_PARTITION_META_RESPONSE" CLEAR_META = "CLEAR_META" CLEAR_META_RESPONSE = "CLEAR_META_RESPONSE" + CLEAR_PARTITION = "CLEAR_PARTITION" + CLEAR_PARTITION_RESPONSE = "CLEAR_PARTITION_RESPONSE" # CHECK_CONSUMPTION CHECK_CONSUMPTION = "CHECK_CONSUMPTION" diff --git a/tutorial/01_core_components.py b/tutorial/01_core_components.py index cd7db6c..3cf3d3e 100644 --- a/tutorial/01_core_components.py +++ b/tutorial/01_core_components.py @@ -174,8 +174,8 @@ def demonstrate_data_workflow(client): print(" ✓ Data matches original!") # Step 5: Clear - print("[Step 5] Clearing partition...") - client.clear(partition_id=partition_id) + print("[Step 5] Clearing partition... (you may also use clear_samples() to clear specific samples)") + client.clear_partition(partition_id=partition_id) print(" ✓ Partition cleared") diff --git a/tutorial/02_metadata_concepts.py b/tutorial/02_metadata_concepts.py index fa2f763..8cc0424 100644 --- a/tutorial/02_metadata_concepts.py +++ b/tutorial/02_metadata_concepts.py @@ -380,7 +380,7 @@ def demonstrate_real_workflow(): print(f" Chunk {i}: Retrieved chunk data: {chunk_data}") # Cleanup - client.clear(partition_id=partition_id) + client.clear_partition(partition_id=partition_id) client.close() ray.shutdown() print("✓ Partition cleared and resources cleaned up") diff --git a/tutorial/03_understanding_controller.py b/tutorial/03_understanding_controller.py index da7c786..fbc4aac 100644 --- a/tutorial/03_understanding_controller.py +++ b/tutorial/03_understanding_controller.py @@ -145,8 +145,8 @@ def demonstrate_partition_isolation(): print(" ✓ Data isolation: 'train' and 'val' partitions are completely independent") # Cleanup - client.clear(partition_id="train") - client.clear(partition_id="val") + client.clear_partition(partition_id="train") + client.clear_partition(partition_id="val") client.close() ray.shutdown() @@ -209,7 +209,7 @@ def demonstrate_dynamic_expansion(): print(" ✓ Columns auto-expand: Can add new fields anytime") # Cleanup - client.clear(partition_id="dynamic") + client.clear_partition(partition_id="dynamic") client.close() ray.shutdown() @@ -256,7 +256,7 @@ def demonstrate_default_consumption_sample_strategy(): print(" ✓ Third get (Task B): samples 0,1") # Cleanup - client.clear(partition_id="sampling") + client.clear_partition(partition_id="sampling") client.close() ray.shutdown() diff --git a/tutorial/04_custom_sampler.py b/tutorial/04_custom_sampler.py index 13120c0..4fcceb3 100644 --- a/tutorial/04_custom_sampler.py +++ b/tutorial/04_custom_sampler.py @@ -245,7 +245,7 @@ def demonstrate_random_sampler_with_replacement(): print(f" ✓ All sampled: {all_sampled}") # Cleanup - client.clear(partition_id="test") + client.clear_partition(partition_id="test") client.close() ray.shutdown() @@ -293,7 +293,7 @@ def demonstrate_random_sampler_without_replacement(): print(f" ✓ Batch 3: {meta3.global_indexes} (none left)") # Cleanup - client.clear(partition_id="test") + client.clear_partition(partition_id="test") client.close() ray.shutdown() @@ -357,7 +357,7 @@ def demonstrate_priority_sampler(): print(f" ✓ Batch 2 high-priority indices: {[i for i in meta2.global_indexes if priority_scores[i] >= 0.1]}") # Cleanup - client.clear(partition_id="test") + client.clear_partition(partition_id="test") client.close() ray.shutdown()