minor update

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_metadata.py

@@ -239,7 +239,7 @@ def test_batch_meta_chunk_by_partition(self):
         assert chunks[3].partition_ids == ["partition_3", "partition_3"]
         assert chunks[3].global_indexes == [13, 17]
 
-        # validate _custom_backend_meta is chunked
+        # validate custom_meta is chunked
         assert 10 in chunks[0].custom_meta
         assert 14 in chunks[0].custom_meta
         assert 18 in chunks[0].custom_meta
@@ -406,22 +406,22 @@ def test_batch_meta_union(self):
 
         batch1 = BatchMeta(
             samples=[
-                SampleMeta(partition_id="partition_0", global_index=0, fields=fields1),
-                SampleMeta(partition_id="partition_0", global_index=1, fields=fields1),
+                SampleMeta(partition_id="partition_0", global_index=8, fields=fields1),
+                SampleMeta(partition_id="partition_0", global_index=9, fields=fields1),
             ],
             _custom_backend_meta={
-                i: {"field1": {"dtype": torch.float32}, "field2": {"dtype": torch.int64}} for i in [0, 1]
+                i: {"field1": {"dtype": torch.float32}, "field2": {"dtype": torch.int64}} for i in [8, 9]
             },
         )
         batch1.extra_info["info1"] = "value1"
 
         batch2 = BatchMeta(
             samples=[
-                SampleMeta(partition_id="partition_0", global_index=0, fields=fields2),
-                SampleMeta(partition_id="partition_0", global_index=1, fields=fields2),
+                SampleMeta(partition_id="partition_0", global_index=8, fields=fields2),
+                SampleMeta(partition_id="partition_0", global_index=9, fields=fields2),
             ],
             _custom_backend_meta={
-                i: {"field2": {"dtype": torch.int64}, "field3": {"dtype": torch.bool}} for i in [0, 1]
+                i: {"field2": {"dtype": torch.int64}, "field3": {"dtype": torch.bool}} for i in [8, 9]
             },
         )
         batch2.extra_info["info2"] = "value2"

transfer_queue/client.py

@@ -314,7 +314,7 @@ async def async_set_custom_meta(
 
         if response_msg.request_type != ZMQRequestType.SET_CUSTOM_META_RESPONSE:
             raise RuntimeError(
-                f"[{self.client_id}]: Failed to set custom metadata from controller {self._controller.id}: "
+                f"[{self.client_id}]: Failed to set custom metadata to controller {self._controller.id}: "
                 f"{response_msg.body.get('message', 'Unknown error')}"
             )
 

transfer_queue/controller.py

@@ -703,8 +703,7 @@ def set_custom_meta(self, custom_meta: dict[int, dict]) -> None:
                         Existing entries will be overwritten.
         """
 
-        for k in custom_meta.keys():
-            self.custom_meta[k] = custom_meta[k]
+        self.custom_meta.update(custom_meta)
 
     # ==================== Statistics and Monitoring ====================
 
@@ -1076,6 +1075,11 @@ def set_custom_meta(self, partition_custom_meta: dict[str, dict[int, dict]]) ->
             partition = self._get_partition(partition_id)
             if partition:
                 partition.set_custom_meta(custom_meta)
+            else:
+                logger.warning(
+                    f"set_custom_meta: partition {partition_id}' not found; "
+                    f"custom_metadata for this partition will be ignored"
+                )
 
     def get_metadata(
         self,

transfer_queue/metadata.py

@@ -215,8 +215,6 @@ def __post_init__(self):
         """Initialize all computed properties during initialization"""
         self.samples = copy.deepcopy(self.samples)
         self.extra_info = copy.deepcopy(self.extra_info)
-        self.custom_meta = copy.deepcopy(self.custom_meta)
-        self._custom_backend_meta = copy.deepcopy(self._custom_backend_meta)
 
         # Basic properties
         object.__setattr__(self, "_size", len(self.samples))
@@ -352,7 +350,7 @@ def update_custom_meta(self, new_meta: dict[int, dict[str, Any]]):
             return
 
         non_exist_global_indexes = set(new_meta.keys()) - set(self.global_indexes)
-        if bool(non_exist_global_indexes):
+        if non_exist_global_indexes:
             raise ValueError(
                 f"Trying to update custom_meta with non-exist global_indexes! {non_exist_global_indexes} "
                 f"do not exist in this batch."
@@ -683,7 +681,7 @@ def union(self, other: "BatchMeta", validate: bool = True) -> Optional["BatchMet
             elif idx in self._custom_backend_meta:
                 merged_custom_backend_meta[idx] = {**self._custom_backend_meta[idx]}
             elif idx in other._custom_backend_meta:
-                merged_custom_backend_meta[idx] = other._custom_backend_meta[idx]
+                merged_custom_backend_meta[idx] = {**other._custom_backend_meta[idx]}
 
         return BatchMeta(
             samples=merged_samples,