Thanks to visit codestin.com
Credit goes to github.com

Skip to content
This repository was archived by the owner on Jan 21, 2026. It is now read-only.

[Misc] Clean up redundant logs & unify log format #160

Merged
0oshowero0 merged 3 commits into from
Dec 24, 2025
Merged

[Misc] Clean up redundant logs & unify log format #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
081e0ba
update
0oshowero0 Dec 24, 2025
86f43db
fix
0oshowero0 Dec 24, 2025
11baaca
fix
0oshowero0 Dec 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions transfer_queue/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()}"
)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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."
)

Expand Down Expand Up @@ -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):
Expand Down
56 changes: 19 additions & 37 deletions transfer_queue/controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 ====================
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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..."
)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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)

Expand All @@ -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(
Expand Down
13 changes: 5 additions & 8 deletions transfer_queue/storage/managers/simple_backend_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()}"
)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions transfer_queue/utils/serial_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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):
Expand Down