[train][data] Fix iter_torch_batches usage of ray.train.torch.get_device when running outside Ray Train (ray-project#57816)

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()`.

---------

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Aydin Abiar <aydin@anyscale.com>

Author: justinvyu

python/ray/data/iterator.py

@@ -408,6 +408,7 @@ def iter_torch_batches(
         """
 
         from ray.train.torch import get_device
+        from ray.train.utils import _in_ray_train_worker
 
         if collate_fn is not None and (dtypes is not None or device != "auto"):
             raise ValueError(
@@ -424,7 +425,7 @@ def iter_torch_batches(
         if device == "auto":
             # Use the appropriate device for Ray Train, or falls back to CPU if
             # Ray Train is not being used.
-            device = get_device()
+            device = get_device() if _in_ray_train_worker() else "cpu"
 
         from ray.air._internal.torch_utils import (
             move_tensors_to_device,

python/ray/train/_internal/session.py

@@ -1197,3 +1197,8 @@ def get_storage() -> StorageContext:
     without notice between minor versions.
     """
     return get_session().storage
+
+
+def _in_ray_train_worker() -> bool:
+    """Check if the current process is a Ray Train V1 worker."""
+    return bool(get_session()) and get_session().world_rank is not None

python/ray/train/tests/test_data_parallel_trainer.py

@@ -14,6 +14,7 @@
 from ray.train.backend import Backend, BackendConfig
 from ray.train.data_parallel_trainer import DataParallelTrainer
 from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.utils import _in_ray_train_worker
 from ray.tune.callback import Callback
 from ray.tune.tune_config import TuneConfig
 from ray.tune.tuner import Tuner
@@ -378,6 +379,16 @@ def train_func():
     trainer.fit()
 
 
+def test_in_ray_train_worker(ray_start_4_cpus):
+    assert not _in_ray_train_worker()
+
+    def train_fn():
+        assert _in_ray_train_worker()
+
+    trainer = DataParallelTrainer(train_fn)
+    trainer.fit()
+
+
 if __name__ == "__main__":
     import sys
 

python/ray/train/tests/test_iter_torch_batches_gpu.py

@@ -281,6 +281,7 @@ def test_custom_batch_collate_fn(
     # Set the device that's returned by device="auto" -> get_device()
     # This is used in `finalize_fn` to move the tensors to the correct device.
     device = torch.device(device)
+    monkeypatch.setattr(ray.train.utils, "_in_ray_train_worker", lambda: True)
     monkeypatch.setattr(ray.train.torch, "get_device", lambda: device)
 
     ds = ray.data.from_items(

python/ray/train/utils.py

@@ -17,3 +17,20 @@ def _log_deprecation_warning(message: str):
         RayDeprecationWarning,
         stacklevel=2,
     )
+
+
+def _in_ray_train_worker() -> bool:
+    from ray.train.v2._internal.constants import is_v2_enabled
+
+    if is_v2_enabled():
+        from ray.train.v2._internal.util import (
+            _in_ray_train_worker as _in_ray_train_v2_worker,
+        )
+
+        return _in_ray_train_v2_worker()
+    else:
+        from ray.train._internal.session import (
+            _in_ray_train_worker as _in_ray_train_v1_worker,
+        )
+
+        return _in_ray_train_v1_worker()

python/ray/train/v2/_internal/util.py

@@ -238,3 +238,14 @@ def construct_user_exception_with_traceback(
     )
     logger.error(f"Error in training function:\n{exc_traceback_str}")
     return UserExceptionWithTraceback(e, traceback_str=exc_traceback_str)
+
+
+def _in_ray_train_worker() -> bool:
+    """Check if the current process is a Ray Train V2 worker."""
+    from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
+
+    try:
+        get_train_fn_utils()
+        return True
+    except RuntimeError:
+        return False

python/ray/train/v2/tests/test_util.py

@@ -1,20 +1,15 @@
 import pytest
 
 import ray
+from ray.train.utils import _in_ray_train_worker
 from ray.train.v2._internal.util import ray_get_safe
-
-
-@pytest.fixture(scope="module")
-def ray_start_4_cpus():
-    ray.init(num_cpus=4)
-    yield
-    ray.shutdown()
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
 
 
 @pytest.mark.parametrize("type", ["task", "actor_task"])
 @pytest.mark.parametrize("failing", [True, False])
 @pytest.mark.parametrize("task_list", [True, False])
-def test_ray_get_safe(type, failing, task_list):
+def test_ray_get_safe(ray_start_4_cpus, type, failing, task_list):
     num_tasks = 4
 
     if type == "task":
@@ -56,6 +51,16 @@ def f(self):
             assert out == 1
 
 
+def test_in_ray_train_worker(ray_start_4_cpus):
+    assert not _in_ray_train_worker()
+
+    def train_fn():
+        assert _in_ray_train_worker()
+
+    trainer = DataParallelTrainer(train_fn)
+    trainer.fit()
+
+
 if __name__ == "__main__":
     import sys