[data] Small (typing) overhaul of valid torch device types

python/ray/data/collate_fn.py

@@ -14,18 +14,18 @@
 
 import numpy as np
 
-from ray.data.block import DataBatch
 from ray.util.annotations import DeveloperAPI
 
 if TYPE_CHECKING:
     import pandas
     import pyarrow
     import torch
 
-    from ray.data.dataset import CollatedData
+    from ray.data.block import DataBatch
+    from ray.data.dataset import CollatedData, TorchDeviceType
 
 
-DataBatchType = TypeVar("DataBatchType", bound=DataBatch)
+DataBatchType = TypeVar("DataBatchType", bound="DataBatch")
 
 TensorSequenceType = Union[
     List["torch.Tensor"],
@@ -226,7 +226,7 @@ class DefaultCollateFn(ArrowBatchCollateFn):
     def __init__(
         self,
         dtypes: Optional[Union["torch.dtype", Dict[str, "torch.dtype"]]] = None,
-        device: Optional[Union[str, "torch.device"]] = None,
+        device: Optional["TorchDeviceType"] = None,
         pin_memory: bool = False,
     ):
         """Initialize the collate function.
@@ -242,7 +242,7 @@ def __init__(
 
         super().__init__()
         self.dtypes = dtypes
-        if isinstance(device, str):
+        if isinstance(device, (str, int)):
             self.device = torch.device(device)
         else:
             self.device = device

python/ray/data/dataset.py

@@ -152,6 +152,12 @@
 CollatedData = TypeVar("CollatedData")
 TorchBatchType = Union[Dict[str, "torch.Tensor"], CollatedData]
 
+TorchDeviceType = Union[str, "torch.device", int]
+"""
+A device identifier, which can be a string (e.g. 'cpu', 'cuda:0'),
+a torch.device object, or other types supported by torch.
+"""
+
 BT_API_GROUP = "Basic Transformations"
 SSR_API_GROUP = "Sorting, Shuffling and Repartitioning"
 SMJ_API_GROUP = "Splitting, Merging, Joining datasets"
@@ -5049,7 +5055,7 @@ def iter_torch_batches(
         prefetch_batches: int = 1,
         batch_size: Optional[int] = 256,
         dtypes: Optional[Union["torch.dtype", Dict[str, "torch.dtype"]]] = None,
-        device: str = "auto",
+        device: Union[TorchDeviceType, Literal["auto"]] = "auto",
         collate_fn: Optional[Callable[[Dict[str, np.ndarray]], CollatedData]] = None,
         drop_last: bool = False,
         local_shuffle_buffer_size: Optional[int] = None,

python/ray/data/iterator.py

@@ -9,6 +9,7 @@
     Iterable,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     TypeVar,
@@ -47,6 +48,7 @@
         Schema,
         TensorFlowTensorBatchType,
         TorchBatchType,
+        TorchDeviceType,
     )
 
 
@@ -272,7 +274,7 @@ def iter_torch_batches(
         prefetch_batches: int = 1,
         batch_size: Optional[int] = 256,
         dtypes: Optional[Union["torch.dtype", Dict[str, "torch.dtype"]]] = None,
-        device: str = "auto",
+        device: Union["TorchDeviceType", Literal["auto"]] = "auto",
         collate_fn: Optional[
             Union[Callable[[Dict[str, np.ndarray]], "CollatedData"], CollateFn]
         ] = None,