fix type comments

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

Author: 0oshowero0

transfer_queue/storage/clients/mooncake_client.py

@@ -141,7 +141,7 @@ def put(self, keys: list[str], values: list[Any]) -> None:
 
         return None
 
-    def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[Tensor]):
+    def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[Tensor]) -> None:
         """Worker thread for putting batch of tensors to MooncakeStore."""
 
         batch_ptrs, batch_sizes, _contiguous_tensors = self._preprocess_tensors_for_put(batch_tensors)
@@ -261,7 +261,7 @@ def _get_bytes_thread_worker(self, batch_keys: list[str], indexes: list[int]) ->
 
         return results, indexes
 
-    def clear(self, keys: list[str], custom_backend_meta=None):
+    def clear(self, keys: list[str], custom_backend_meta: Optional[list[Any]] = None) -> None:
         """Deletes multiple keys from MooncakeStore.
 
         Args:
@@ -280,10 +280,10 @@ def close(self):
             self._store = None
 
     @staticmethod
-    def _preprocess_tensors_for_put(values: list[Tensor]) -> tuple[list[Any], list[Any], list[Tensor]]:
-        ptr_list = []
-        size_list = []
-        tensor_list = []  # hold reference for the contiguous tensor
+    def _preprocess_tensors_for_put(values: list[Tensor]) -> tuple[list[int], list[int], list[Tensor]]:
+        ptr_list: list[int] = []
+        size_list: list[int] = []
+        tensor_list: list[Tensor] = []  # hold reference for the contiguous tensor
         for t in values:
             # TODO: support gpu direct rdma and use different data paths.
             #       For GPU, it's more reasonable to perform data copy since

transfer_queue/storage/clients/yuanrong_client.py

@@ -57,7 +57,7 @@ def supports_put(self, value: Any) -> bool:
         """Check if this strategy can store the given value."""
 
     @abstractmethod
-    def put(self, keys: list[str], values: list[Any]):
+    def put(self, keys: list[str], values: list[Any]) -> None:
         """Store key-value pairs using this strategy."""
 
     @abstractmethod
@@ -73,7 +73,7 @@ def supports_clear(self, strategy_tag: Any) -> bool:
         """Check if this strategy owns data identified by metadata."""
 
     @abstractmethod
-    def clear(self, keys: list[str]):
+    def clear(self, keys: list[str]) -> None:
         """Delete keys from storage."""
 
 
@@ -131,7 +131,7 @@ def supports_put(self, value: Any) -> bool:
         # Only contiguous NPU tensors are supported by this adapter.
         return value.is_contiguous()
 
-    def put(self, keys: list[str], values: list[Any]):
+    def put(self, keys: list[str], values: list[Any]) -> None:
         """Store NPU tensors in batches; deletes before overwrite."""
         for i in range(0, len(keys), self.KEYS_LIMIT):
             batch_keys = keys[i : i + self.KEYS_LIMIT]
@@ -169,22 +169,22 @@ def supports_clear(self, strategy_tag: str) -> bool:
         """Matches 'DsTensorClient' strategy tag."""
         return isinstance(strategy_tag, str) and strategy_tag == self.strategy_tag()
 
-    def clear(self, keys: list[str]):
+    def clear(self, keys: list[str]) -> None:
         """Delete NPU tensor keys in batches."""
         for i in range(0, len(keys), self.KEYS_LIMIT):
             batch = keys[i : i + self.KEYS_LIMIT]
             # Todo(dpj): Test call clear when no (key,value) put in ds
             self._ds_client.delete(batch)
 
-    def _create_empty_npu_tensorlist(self, shapes: list, dtypes: list):
+    def _create_empty_npu_tensorlist(self, shapes: list[Any], dtypes: list[Any]) -> list[Tensor]:
         """
         Create a list of empty NPU tensors with given shapes and dtypes.
 
         Args:
             shapes (list): List of tensor shapes (e.g., [(3,), (2, 4)])
             dtypes (list): List of torch dtypes (e.g., [torch.float32, torch.int64])
         Returns:
-            list: List of uninitialized NPU tensors
+            list[Tensor]: List of uninitialized NPU tensors
         """
         tensors: list[Tensor] = []
         for shape, dtype in zip(shapes, dtypes, strict=True):
@@ -243,7 +243,7 @@ def supports_put(self, value: Any) -> bool:
         """Accepts any Python object."""
         return True
 
-    def put(self, keys: list[str], values: list[Any]):
+    def put(self, keys: list[str], values: list[Any]) -> None:
         """Store objects via zero-copy serialization in batches."""
         for i in range(0, len(keys), self.PUT_KEYS_LIMIT):
             batch_keys = keys[i : i + self.PUT_KEYS_LIMIT]
@@ -267,7 +267,7 @@ def supports_clear(self, strategy_tag: str) -> bool:
         """Matches 'KVClient' strategy tag."""
         return isinstance(strategy_tag, str) and strategy_tag == self.strategy_tag()
 
-    def clear(self, keys: list[str]):
+    def clear(self, keys: list[str]) -> None:
         """Delete keys in batches."""
         for i in range(0, len(keys), self.GET_CLEAR_KEYS_LIMIT):
             batch_keys = keys[i : i + self.GET_CLEAR_KEYS_LIMIT]
@@ -433,7 +433,13 @@ def put_task(strategy, indexes):
                 strategy_tags[original_index] = tag
         return strategy_tags
 
-    def get(self, keys: list[str], shapes=None, dtypes=None, custom_backend_meta=None) -> list[Any]:
+    def get(
+        self,
+        keys: list[str],
+        shapes: Optional[list[Any]] = None,
+        dtypes: Optional[list[Any]] = None,
+        custom_backend_meta: Optional[list[str]] = None,
+    ) -> list[Any]:
         """Retrieves multiple values from remote storage with expected metadata.
 
         Requires shape and dtype hints to reconstruct NPU tensors correctly.
@@ -472,7 +478,7 @@ def get_task(strategy, indexes):
                 results[original_index] = value
         return results
 
-    def clear(self, keys: list[str], custom_backend_meta=None):
+    def clear(self, keys: list[str], custom_backend_meta: Optional[list[str]] = None) -> None:
         """Deletes multiple keys from remote storage.
 
         Args:
@@ -513,8 +519,8 @@ def _route_to_strategies(
                    The order must correspond to the original keys.
             selector: A function that determines whether a strategy supports an item.
                      Signature: `(strategy: StorageStrategy, item: Any) -> bool`.
-            failback: If True, items that don't match any strategy will be ignored (not included in output).
-                      If False, a ValueError will be raised for any unmatched item.
+            ignore_unmatched: If True, items that don't match any strategy will be ignored (not included in output).
+                              If False, a ValueError will be raised for any unmatched item.
 
         Returns:
             A dictionary mapping each active strategy to a list of indexes in `items`

transfer_queue/storage/managers/base.py

@@ -401,19 +401,20 @@ def _generate_keys(field_names: list[str], global_indexes: list[int]) -> list[st
         return [pfx + sfx for sfx, pfx in itertools.product(keys_suffixes, keys_prefixes)]
 
     @staticmethod
-    def _generate_values(data: TensorDict) -> list[Tensor]:
+    def _generate_values(data: TensorDict) -> list[Any]:
         """
-        Extract and flatten tensor values from a TensorDict in field-major order.
+        Extract and flatten values from a TensorDict in field-major order.
         Values are ordered by sorted field names, then by row (sample) order within each field.
         This matches the key order generated by `_generate_keys`.
 
         Args:
-            data (TensorDict): Input data where keys are field names and values are tensors.
+            data (TensorDict): Input data where keys are field names and values are tensors or any type
+                               wrapped by NonTensorStack.
         Returns:
-            list[Tensor]: Flattened list of tensors, e.g.,
-                          [data[field_a][0], data[field_a][1], data[field_a][2], ..., data[field_b][0], ...]
+            list[Any]: Flattened list of values, e.g.,
+                       [data[field_a][0], data[field_a][1], data[field_a][2], ..., data[field_b][0], ...]
         """
-        results: list[Tensor] = []
+        results: list[Any] = []
         for field in sorted(data.keys()):
             field_data = data[field]
             if isinstance(field_data, Tensor) and field_data.is_nested:
@@ -457,17 +458,17 @@ def _get_executor(self) -> ThreadPoolExecutor:
         assert self._multi_threads_executor is not None
         return self._multi_threads_executor
 
-    def _merge_tensors_to_tensordict(self, metadata: BatchMeta, values: list[Tensor]) -> TensorDict:
+    def _merge_tensors_to_tensordict(self, metadata: BatchMeta, values: list[Any]) -> TensorDict:
         """
         Reconstruct a TensorDict from a list of values using metadata.
         The values list is assumed to be in the same order as keys generated by `_generate_keys`.
         According to field names and global indexes in metadata, this method can determine
-        which dict key and which row this tensor belongs to. Then it reshapes the flat tensors list
+        which dict key and which row this value belongs to. Then it reshapes the flat values list
         back into a structured TensorDict .
 
         Args:
             metadata (BatchMeta): Metadata containing global indexes and field names.
-            values (list[Tensor]): List of tensors in field-major order.
+            values (list[Any]): List of values in field-major order.
         Returns:
             TensorDict: Reconstructed tensor dictionary with batch size equal to number of samples.
         """
@@ -534,7 +535,9 @@ def process_field(field_idx: int):
         return TensorDict(merged_data, batch_size=num_samples)
 
     @staticmethod
-    def _get_shape_type_custom_backend_meta_list(metadata: BatchMeta):
+    def _get_shape_type_custom_backend_meta_list(
+        metadata: BatchMeta,
+    ) -> tuple[list[torch.Size], list[torch.dtype], list[Any]]:
         """
         Extract the expected shape, dtype, and custom_backend_meta for each field-sample pair in metadata.
         The order matches the key/value order: sorted by field name, then by global index.