set target-verion py310 and enable modern type annotation rules (Optional/Union → X | None)

Signed-off-by: ji-huazhong <hzji210@gmail.com>

Author: ji-huazhong

pyproject.toml

@@ -31,6 +31,7 @@ requires-python = ">=3.10"
 # Note: While the formatter will attempt to format lines such that they remain within the line-length,
 # it isn't a hard upper bound, and formatted lines may exceed the line-length.
 line-length = 120
+target-version = "py310"
 
 [tool.ruff.lint]
 isort = {known-first-party = ["transfer_queue"]}
@@ -60,7 +61,7 @@ ignore = [
     # `.log()` statement uses f-string
     "G004",
     # X | None for type annotations
-    "UP045",
+    # "UP045",
     # deprecated import
     "UP035",
 ]

transfer_queue/client.py

@@ -16,7 +16,7 @@
 import asyncio
 import os
 import threading
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 import torch
 import zmq
@@ -104,9 +104,9 @@ async def async_get_meta(
         batch_size: int,
         partition_id: str,
         mode: str = "fetch",
-        task_name: Optional[str] = None,
-        sampling_config: Optional[dict[str, Any]] = None,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        task_name: str | None = None,
+        sampling_config: dict[str, Any] | None = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
         """Asynchronously fetch data metadata from the controller via ZMQ.
 
@@ -191,7 +191,7 @@ async def async_get_meta(
     async def async_set_custom_meta(
         self,
         metadata: BatchMeta,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> None:
         """
         Asynchronously send custom metadata to the controller.
@@ -264,9 +264,9 @@ async def async_set_custom_meta(
     async def async_put(
         self,
         data: TensorDict,
-        metadata: Optional[BatchMeta] = None,
-        partition_id: Optional[str] = None,
-        data_parser: Optional[Callable[[Any], Any]] = None,
+        metadata: BatchMeta | None = None,
+        partition_id: str | None = None,
+        data_parser: Callable[[Any], Any] | None = None,
     ) -> BatchMeta:
         """Asynchronously write data to storage units based on metadata.
 
@@ -575,8 +575,8 @@ async def async_get_consumption_status(
         self,
         task_name: str,
         partition_id: str,
-        socket: Optional[zmq.asyncio.Socket] = None,
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+        socket: zmq.asyncio.Socket | None = None,
+    ) -> tuple[Tensor | None, Tensor | None]:
         """Get consumption status for current partition in a specific task.
 
         Args:
@@ -638,8 +638,8 @@ async def async_get_production_status(
         self,
         data_fields: list[str],
         partition_id: str,
-        socket: Optional[zmq.asyncio.Socket] = None,
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+        socket: zmq.asyncio.Socket | None = None,
+    ) -> tuple[Tensor | None, Tensor | None]:
         """Get production status for specific data fields and partition.
 
         Args:
@@ -769,8 +769,8 @@ async def async_check_production_status(
     async def async_reset_consumption(
         self,
         partition_id: str,
-        task_name: Optional[str] = None,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        task_name: str | None = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> bool:
         """Asynchronously reset consumption status for a partition.
 
@@ -830,7 +830,7 @@ async def async_reset_consumption(
     @with_controller_socket
     async def async_get_partition_list(
         self,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> list[str]:
         """Asynchronously fetch the list of partition ids from the controller.
 
@@ -879,7 +879,7 @@ async def async_kv_retrieve_meta(
         keys: list[str] | str,
         partition_id: str,
         create: bool = False,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
         """Asynchronously retrieve BatchMeta from the controller using user-specified keys.
 
@@ -944,7 +944,7 @@ async def async_kv_retrieve_keys(
         self,
         global_indexes: list[int] | int,
         partition_id: str,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> list[str]:
         """Asynchronously retrieve keys according to global_indexes from the controller.
 
@@ -1005,8 +1005,8 @@ async def async_kv_retrieve_keys(
     @with_controller_socket
     async def async_kv_list(
         self,
-        partition_id: Optional[str] = None,
-        socket: Optional[zmq.asyncio.Socket] = None,
+        partition_id: str | None = None,
+        socket: zmq.asyncio.Socket | None = None,
     ) -> dict[str, dict[str, Any]]:
         """Asynchronously retrieve keys and custom_meta from the controller for one or all partitions.
 
@@ -1145,8 +1145,8 @@ def get_meta(
         batch_size: int,
         partition_id: str,
         mode: str = "fetch",
-        task_name: Optional[str] = None,
-        sampling_config: Optional[dict[str, Any]] = None,
+        task_name: str | None = None,
+        sampling_config: dict[str, Any] | None = None,
     ) -> BatchMeta:
         """Synchronously fetch data metadata from the controller via ZMQ.
 
@@ -1234,9 +1234,9 @@ def set_custom_meta(self, metadata: BatchMeta) -> None:
     def put(
         self,
         data: TensorDict,
-        metadata: Optional[BatchMeta] = None,
-        partition_id: Optional[str] = None,
-        data_parser: Optional[Callable[[Any], Any]] = None,
+        metadata: BatchMeta | None = None,
+        partition_id: str | None = None,
+        data_parser: Callable[[Any], Any] | None = None,
     ) -> BatchMeta:
         """Synchronously write data to storage units based on metadata.
 
@@ -1356,7 +1356,7 @@ def get_consumption_status(
         self,
         task_name: str,
         partition_id: str,
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+    ) -> tuple[Tensor | None, Tensor | None]:
         """Synchronously get consumption status for a specific task and partition.
 
         Args:
@@ -1384,7 +1384,7 @@ def get_production_status(
         self,
         data_fields: list[str],
         partition_id: str,
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+    ) -> tuple[Tensor | None, Tensor | None]:
         """Synchronously get production status for specific data fields and partition.
 
         Args:
@@ -1454,7 +1454,7 @@ def check_production_status(self, data_fields: list[str], partition_id: str) ->
         """
         return self._check_production_status(data_fields=data_fields, partition_id=partition_id)
 
-    def reset_consumption(self, partition_id: str, task_name: Optional[str] = None) -> bool:
+    def reset_consumption(self, partition_id: str, task_name: str | None = None) -> bool:
         """Synchronously reset consumption status for a partition.
 
         This allows the same data to be re-consumed, useful for debugging scenarios
@@ -1540,7 +1540,7 @@ def kv_retrieve_keys(
 
     def kv_list(
         self,
-        partition_id: Optional[str] = None,
+        partition_id: str | None = None,
     ) -> dict[str, dict[str, Any]]:
         """Synchronously retrieve keys and custom_meta from the controller for one or all partitions.
 

transfer_queue/controller.py

@@ -21,7 +21,7 @@
 from itertools import groupby
 from operator import itemgetter
 from threading import Lock, Thread
-from typing import Any, Optional
+from typing import Any
 from uuid import uuid4
 
 import numpy as np
@@ -195,10 +195,10 @@ class FieldMeta:
     """
 
     global_indexes: set[int] = field(default_factory=set)
-    dtype: Optional[Any] = None
-    shape: Optional[tuple] = None  # None when is_nested=True
-    is_nested: Optional[bool] = None
-    is_non_tensor: Optional[bool] = None
+    dtype: Any | None = None
+    shape: tuple | None = None  # None when is_nested=True
+    is_nested: bool | None = None
+    is_non_tensor: bool | None = None
 
     per_sample_shapes: dict[int, tuple] = field(default_factory=dict)  # {global_idx: shape}
 
@@ -495,7 +495,7 @@ def update_production_status(
         global_indices: list[int],
         field_names: list[str],
         field_schema: dict[str, dict[str, Any]],
-        custom_backend_meta: Optional[dict[int, dict[str, Any]]] = None,
+        custom_backend_meta: dict[int, dict[str, Any]] | None = None,
     ) -> bool:
         """
         Update production status for specific samples and fields.
@@ -560,7 +560,7 @@ def _update_field_metadata(
         self,
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
-        custom_backend_meta: Optional[dict[int, dict[str, Any]]] = None,
+        custom_backend_meta: dict[int, dict[str, Any]] | None = None,
     ):
         """Update field metadata from columnar field_schema."""
         if not global_indexes:
@@ -645,7 +645,7 @@ def get_consumption_status(self, task_name: str, mask: bool = False) -> tuple[Te
             consumption_status = self.consumption_status[task_name]
         return partition_global_index, consumption_status
 
-    def reset_consumption(self, task_name: Optional[str] = None):
+    def reset_consumption(self, task_name: str | None = None):
         """
         Reset consumption status for a specific task or all tasks.
 
@@ -670,7 +670,7 @@ def reset_consumption(self, task_name: Optional[str] = None):
     # ==================== Production Status Interface ====================
     def get_production_status_for_fields(
         self, field_names: list[str], mask: bool = False
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+    ) -> tuple[Tensor | None, Tensor | None]:
         """
         Check if all samples for specified fields are fully produced and ready.
 
@@ -1049,7 +1049,7 @@ def create_partition(self, partition_id: str) -> bool:
         logger.info(f"Created partition {partition_id} with {TQ_PRE_ALLOC_SAMPLE_NUM} pre-allocated indexes")
         return True
 
-    def _get_partition(self, partition_id: str) -> Optional[DataPartitionStatus]:
+    def _get_partition(self, partition_id: str) -> DataPartitionStatus | None:
         """
         Get partition status information.
 
@@ -1061,7 +1061,7 @@ def _get_partition(self, partition_id: str) -> Optional[DataPartitionStatus]:
         """
         return self.partitions.get(partition_id)
 
-    def get_partition_snapshot(self, partition_id: str) -> Optional[DataPartitionStatus]:
+    def get_partition_snapshot(self, partition_id: str) -> DataPartitionStatus | None:
         """
         Get a copy of partition status information, without threading.Lock().
 
@@ -1109,7 +1109,7 @@ def update_production_status(
         partition_id: str,
         global_indexes: list[int],
         field_schema: dict[str, dict[str, Any]],
-        custom_backend_meta: Optional[dict[int, dict[str, Any]]] = None,
+        custom_backend_meta: dict[int, dict[str, Any]] | None = None,
     ) -> bool:
         """
         Update production status for specific samples and fields in a partition.
@@ -1139,7 +1139,7 @@ def update_production_status(
         return success
 
     # ==================== Data Consumption API ====================
-    def get_consumption_status(self, partition_id: str, task_name: str) -> tuple[Optional[Tensor], Optional[Tensor]]:
+    def get_consumption_status(self, partition_id: str, task_name: str) -> tuple[Tensor | None, Tensor | None]:
         """
         Get or create consumption status for a specific task and partition.
         Delegates to the partition's own method.
@@ -1159,9 +1159,7 @@ def get_consumption_status(self, partition_id: str, task_name: str) -> tuple[Opt
 
         return partition.get_consumption_status(task_name, mask=True)
 
-    def get_production_status(
-        self, partition_id: str, data_fields: list[str]
-    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
+    def get_production_status(self, partition_id: str, data_fields: list[str]) -> tuple[Tensor | None, Tensor | None]:
         """
         Check if all samples for specified fields are fully produced in a partition.
 
@@ -1226,7 +1224,7 @@ def get_metadata(
         mode: str = "fetch",
         task_name: str | None = None,
         batch_size: int | None = None,
-        sampling_config: Optional[dict[str, Any]] = None,
+        sampling_config: dict[str, Any] | None = None,
         *args,
         **kwargs,
     ) -> BatchMeta:
@@ -1494,7 +1492,7 @@ def clear_partition(self, partition_id: str, clear_consumption: bool = True):
         self.partitions.pop(partition_id)
         self.sampler.clear_cache(partition_id)
 
-    def reset_consumption(self, partition_id: str, task_name: Optional[str] = None):
+    def reset_consumption(self, partition_id: str, task_name: str | None = None):
         """
         Reset consumption status for a partition without clearing the actual data.
 
@@ -1641,7 +1639,7 @@ def kv_retrieve_keys(
         self,
         global_indexes: list[int],
         partition_id: str,
-    ) -> list[Optional[str]]:
+    ) -> list[str | None]:
         """
         Retrieve keys from the controller using a list of global_indexes.
 

transfer_queue/dataloader/streaming_dataloader.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
 
 import torch
 from tensordict import TensorDict
@@ -80,7 +79,7 @@ def __init__(
         pin_memory: bool = False,
         worker_init_fn=None,
         multiprocessing_context=None,
-        prefetch_factor: Optional[int] = None,
+        prefetch_factor: int | None = None,
         persistent_workers: bool = False,
         pin_memory_device: str = "",
     ):