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

## Why are these changes needed?

Follow up of RLlib PR: #55291
Torch `device` allows input types of `str | int | torch.device`, this PR
unifies the type in a type variable and allows for the `int` type as
well.

Upstream air PR: #56745

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ x I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [x] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] This PR is not tested :(

---------

Signed-off-by: Daniel Sperber <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: Daraan

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 an integer (e.g. 0 for 'cuda:0').
+"""
+
 BT_API_GROUP = "Basic Transformations"
 SSR_API_GROUP = "Sorting, Shuffling and Repartitioning"
 SMJ_API_GROUP = "Splitting, Merging, Joining datasets"
@@ -5318,7 +5324,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,