[train][data] Fix iter_torch_batches usage of ray.train.torch.get_device when running outside Ray Train

[justinvyu]

Description

Train V2 doesn't allow running ray.train.torch.get_device outside of a Ray Train worker spawned by a trainer.fit() call. Previously, get_device() returns the 0th index GPU of the ray.get_gpu_ids() assigned to the current process, or "cpu" if the current process wasn't assigned GPUs via ray.remote(num_gpus=x).

This PR introduces a utility to detect whether we're running inside a Ray Train worker process or not (in v1 and v2) and updates Ray Data's iter_torch_batches to only call get_device() if in a Train worker process.

This introduces a slight API breakage for users who spawned a custom GPU Ray task and used iter_torch_batches or get_device().

[gemini-code-assist]

Code Review

This pull request introduces a utility _in_ray_train_worker to correctly detect whether code is running within a Ray Train worker, for both V1 and V2. This utility is then used in iter_torch_batches to conditionally call ray.train.torch.get_device, fixing an issue where get_device would fail when used outside a Ray Train context with Train V2. The changes are well-structured, with separate logic for V1 and V2, and are accompanied by good test coverage. My feedback includes a couple of minor suggestions to improve code style and maintainability.

[srinathk10]

nit: Would it be possible roll this condition into get_device?

[justinvyu]

I considered this, but our current stance is that we don't want people calling get_device() (Train worker utils) outside of Train workers (it will error in V2), so I fully gate the call to this method.