diff --git a/transfer_queue/client.py b/transfer_queue/client.py index b40e0b7..496bcbe 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -129,7 +129,7 @@ async def wrapper(self, *args, **kwargs): try: sock.connect(address) - logger.info( + logger.debug( f"[{self.client_id}]: Connected to Controller {server_info.id} at {address} " f"with identity {identity.decode()}" ) @@ -237,8 +237,7 @@ async def async_get_meta( response_serialized = await socket.recv_multipart() response_msg = ZMQMessage.deserialize(response_serialized) logger.debug( - f"[{self.client_id}]: Client get datameta response: {response_msg} " - f"from controller {self._controller.id}" + f"[{self.client_id}]: Client get_meta response: {response_msg} from controller {self._controller.id}" ) if response_msg.request_type == ZMQRequestType.GET_META_RESPONSE: @@ -331,14 +330,13 @@ async def async_put( if not metadata or metadata.size == 0: raise ValueError("metadata cannot be none or empty") - logger.debug(f"[{self.client_id}]: Put data with data: {data}") with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_put" ): await self.storage_manager.put_data(data, metadata) - logger.info( + logger.debug( f"[{self.client_id}]: partition {partition_id} put {metadata.size} samples to storage units successfully." ) @@ -386,6 +384,8 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: ): results = await self.storage_manager.get_data(metadata) + logger.debug(f"[{self.client_id}]: get_data with {metadata.size} samples successfully.") + return results async def async_clear_partition(self, partition_id: str): diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index e0c9bd6..bbea392 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -284,11 +284,11 @@ def ensure_samples_capacity(self, required_samples: int) -> bool: self.consumption_status[task_name] = expanded_consumption logger.debug( - f"Expanded partition {self.partition_id} from {current_sample_space} to {new_samples} samples " - f"(added {min_expansion} samples)" + f"Expanded partition {self.partition_id} from {current_sample_space} " + f"to {new_samples} samples (added {min_expansion} samples)" ) - def ensure_fields_capacity(self, required_fields: int) -> bool: + def ensure_fields_capacity(self, required_fields: int): """ Ensure the production status tensor has enough columns for the required fields. Dynamically expands if needed using unified minimum expansion size. @@ -313,8 +313,8 @@ def ensure_fields_capacity(self, required_fields: int) -> bool: self.production_status = expanded_tensor logger.debug( - f"Expanded partition {self.partition_id} from {current_fields} to {new_fields} fields " - f"(added {min_expansion} fields)" + f"Expanded partition {self.partition_id} from {current_fields} " + f"to {new_fields} fields (added {min_expansion} fields)" ) # ==================== Production Status Interface ==================== @@ -766,22 +766,6 @@ def list_partitions(self) -> list[str]: """ return list(self.partitions.keys()) - def delete_partition(self, partition_id: str) -> bool: - """ - Delete a partition and all its data. - - Args: - partition_id: ID of the partition to delete - - Returns: - True if partition was deleted, False if it didn't exist - """ - if partition_id in self.partitions: - del self.partitions[partition_id] - logger.info(f"Deleted partition {partition_id}") - return True - return False - # ==================== Partition Index Management API ==================== def get_partition_index_range(self, partition: DataPartitionStatus) -> set: @@ -829,8 +813,8 @@ def update_production_status( success = partition.update_production_status(global_indexes, field_names, dtypes, shapes) if success: logger.debug( - f"Updated production status for partition {partition_id}: samples={global_indexes}, " - f"fields={field_names}" + f"[{self.controller_id}]: Updated production status for partition {partition_id}: " + f"samples={global_indexes}, fields={field_names}" ) return success @@ -929,9 +913,9 @@ def get_metadata( if len(ready_for_consume_indexes) < batch_size: if self.polling_mode: logger.debug( - f"Not enough data for task {task_name} in partition {partition_id}. " - f"Required: {batch_size}, Available: {len(ready_for_consume_indexes)}. " - f"Returning None due to polling mode." + f"[{self.controller_id}]: Not enough data for task {task_name} in partition {partition_id}." + f" Required: {batch_size}, Available: {len(ready_for_consume_indexes)}." + f" Returning None due to polling mode." ) return BatchMeta.empty() if time.time() - start_time > TQ_CONTROLLER_GET_METADATA_TIMEOUT: @@ -940,8 +924,8 @@ def get_metadata( f"Required: {batch_size}, Available: {len(ready_for_consume_indexes)}" ) logger.warning( - f"Insufficient data for task {task_name}. Required: {batch_size} samples with " - f"fields {data_fields} in partition {partition_id}, but only have " + f"[{self.controller_id}]: Insufficient data for task {task_name}. Required: {batch_size} " + f"samples with fields {data_fields} in partition {partition_id}, but only have " f"{len(ready_for_consume_indexes)} samples meeting the criteria. " f"Retrying in {TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL}s..." ) @@ -978,8 +962,6 @@ def get_metadata( partition = self.partitions[partition_id] partition.mark_consumed(task_name, consumed_indexes) - logger.debug(f"get_metadata: {metadata}") - return metadata def scan_data_status( @@ -1098,7 +1080,7 @@ def clear_partition(self, partition_id: str, clear_consumption: bool = True): clear_consumption: Whether to also clear consumption status """ - logger.debug(f"Clearing data for partition {partition_id}") + logger.debug(f"[{self.controller_id}]: clearing metadata in partition {partition_id}") partition = self._get_partition(partition_id) if not partition: @@ -1120,7 +1102,7 @@ def clear_meta(self, global_indexes: list[int], partition_ids: list[str], clear_ """ logger.debug( - f"{self.controller_id}: Clearing meta with global_indexes {global_indexes} in partition {partition_ids}" + 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: @@ -1223,18 +1205,18 @@ def _wait_connection(self): self._connected_storage_managers.add(storage_manager_id) storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown") logger.debug( - f"Controller {self.controller_id} received handshake from " + f"[{self.controller_id}]: received handshake from " f"storage manager {storage_manager_id} (type: {storage_manager_type}). " f"Total connected: {len(self._connected_storage_managers)}" ) else: logger.debug( - f"Controller {self.controller_id} received duplicate handshake from " + f"[{self.controller_id}]: received duplicate handshake from " f"storage manager {storage_manager_id}. Resending ACK." ) except Exception as e: - logger.error(f"Controller {self.controller_id} error processing handshake: {e}") + logger.error(f"[{self.controller_id}]: error processing handshake: {e}") def _start_process_handshake(self): """Start the handshake process thread.""" @@ -1395,7 +1377,7 @@ def _process_request(self): def _update_data_status(self): """Process data status update messages from storage units - adapted for partitions.""" - logger.info(f"[{self.controller_id}]: start receiving update_data_status requests...") + logger.debug(f"[{self.controller_id}]: start receiving update_data_status requests...") perf_monitor = IntervalPerfMonitor(caller_name=self.controller_id) @@ -1420,7 +1402,7 @@ def _update_data_status(self): ) if success: - logger.info(f"Updated production status for partition {partition_id}") + logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}") # Send acknowledgment response_msg = ZMQMessage.create( diff --git a/transfer_queue/storage/managers/simple_backend_manager.py b/transfer_queue/storage/managers/simple_backend_manager.py index 4967e1d..f361655 100644 --- a/transfer_queue/storage/managers/simple_backend_manager.py +++ b/transfer_queue/storage/managers/simple_backend_manager.py @@ -142,7 +142,7 @@ async def wrapper(self, *args, **kwargs): # Timeouts to avoid indefinite await on recv/send sock.setsockopt(zmq.RCVTIMEO, TQ_SIMPLE_STORAGE_MANAGER_RECV_TIMEOUT * 1000) sock.setsockopt(zmq.SNDTIMEO, TQ_SIMPLE_STORAGE_MANAGER_SEND_TIMEOUT * 1000) - logger.info( + logger.debug( f"[{self.storage_manager_id}]: Connected to StorageUnit {server_info.id} at {address} " f"with identity {identity.decode()}" ) @@ -178,7 +178,7 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None: metadata: BatchMeta containing storage location information. """ - logger.info(f"{__class__.__name__}: receive put_data request, putting {metadata.size} samples.") + logger.debug(f"[{self.storage_manager_id}]: receive put_data request, putting {metadata.size} samples.") # group samples by storage unit storage_meta_groups = build_storage_meta_groups( @@ -277,7 +277,7 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: TensorDict containing the retrieved data. """ - logger.info(f"{__class__.__name__}: receive get_data request, getting {metadata.size} samples.") + logger.debug(f"[{self.storage_manager_id}]: receive get_data request, getting {metadata.size} samples.") # group samples by storage unit storage_meta_groups = build_storage_meta_groups( @@ -347,10 +347,6 @@ async def _get_from_single_storage_unit( await socket.send_multipart(request_msg.serialize()) messages = await socket.recv_multipart() response_msg = ZMQMessage.deserialize(messages) - logger.info( - f"[{self.storage_manager_id}]: get data response from storage unit " - f"{target_storage_unit}: {response_msg}" - ) if response_msg.request_type == ZMQRequestType.GET_DATA_RESPONSE: # Return data and index information from this storage unit @@ -371,6 +367,8 @@ async def clear_data(self, metadata: BatchMeta) -> None: metadata: BatchMeta that contains metadata for data clearing. """ + logger.debug(f"[{self.storage_manager_id}]: receive clear_data request, clearing {metadata.size} samples.") + # group samples by storage unit storage_meta_groups = build_storage_meta_groups( metadata, self.global_index_storage_unit_mapping, self.global_index_local_index_mapping @@ -408,7 +406,6 @@ async def _clear_single_storage_unit(self, local_indexes, target_storage_unit=No f"{response_msg.body.get('message', 'Unknown error')}" ) - logger.info(f"[{self.storage_manager_id}]: Successfully clear storage unit {target_storage_unit}") except Exception as e: logger.error(f"[{self.storage_manager_id}]: Error clearing storage unit {target_storage_unit}: {str(e)}") raise diff --git a/transfer_queue/utils/serial_utils.py b/transfer_queue/utils/serial_utils.py index 6db9e40..a358d09 100644 --- a/transfer_queue/utils/serial_utils.py +++ b/transfer_queue/utils/serial_utils.py @@ -197,7 +197,7 @@ def serialization(obj: Any) -> list[bytestr]: through pickle. """ - logger.info(f"Serializing an obj with TQ_ZERO_COPY_SERIALIZATION={TQ_ZERO_COPY_SERIALIZATION}") + logger.debug(f"Serializing an obj with TQ_ZERO_COPY_SERIALIZATION={TQ_ZERO_COPY_SERIALIZATION}") if TQ_ZERO_COPY_SERIALIZATION: pickled_bytes, tensors = _internal_rpc_pickler.serialize(obj) @@ -219,7 +219,7 @@ def serialization(obj: Any) -> list[bytestr]: def deserialization(data: list[bytestr] | bytestr) -> Any: """Deserialize any object from serialized data.""" - logger.info(f"Deserializing an obj with TQ_ZERO_COPY_SERIALIZATION={TQ_ZERO_COPY_SERIALIZATION}") + logger.debug(f"Deserializing an obj with TQ_ZERO_COPY_SERIALIZATION={TQ_ZERO_COPY_SERIALIZATION}") if TQ_ZERO_COPY_SERIALIZATION: if isinstance(data, list):