[RAY AIR] set the correct gpu id in TorchTrainer (ray-project#26493)

Signed-off-by: Xiaowei Jiang <xwjiang2010@gmail.com>

python/ray/train/BUILD

@@ -315,7 +315,7 @@ py_test(
 
 py_test(
     name = "test_tune",
-    size = "medium",
+    size = "large",
     srcs = ["tests/test_tune.py"],
     tags = ["team:ml", "exclusive", "tune"],
     deps = [":train_lib"]

python/ray/train/tests/test_gpu.py

@@ -1,14 +1,21 @@
 import os
 from timeit import default_timer as timer
+from collections import Counter
 
+from unittest.mock import patch
 import pytest
 import torch
 import torchvision
-from test_tune import torch_fashion_mnist, tune_tensorflow_mnist
+from test_tune import (
+    torch_fashion_mnist,
+    tune_tensorflow_mnist,
+)
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 
 import ray
+from ray.cluster_utils import Cluster
+
 import ray.train as train
 from ray.train import Trainer, TrainingCallback
 from ray.air.config import ScalingConfig
@@ -25,6 +32,7 @@
 from ray.train.examples.torch_linear_example import LinearDataset
 from ray.train.horovod.horovod_trainer import HorovodTrainer
 from ray.train.tensorflow.tensorflow_trainer import TensorflowTrainer
+from ray.train.torch import TorchConfig
 from ray.train.torch.torch_trainer import TorchTrainer
 
 
@@ -43,6 +51,20 @@ def ray_start_1_cpu_1_gpu():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_2_node_4_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=8, num_gpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
 # TODO: Refactor as a backend test.
 @pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1])
 def test_torch_get_device(ray_start_4_cpus_2_gpus, num_gpus_per_worker):
@@ -70,6 +92,40 @@ def train_fn():
         )
 
 
+# TODO: Refactor as a backend test.
+@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
+def test_torch_get_device_dist(ray_2_node_4_gpu, num_gpus_per_worker):
+    @patch("torch.cuda.is_available", lambda: True)
+    def train_fn():
+        return train.torch.get_device().index
+
+    trainer = Trainer(
+        TorchConfig(backend="gloo"),
+        num_workers=int(8 / num_gpus_per_worker),
+        use_gpu=True,
+        resources_per_worker={"GPU": num_gpus_per_worker},
+    )
+    trainer.start()
+    devices = trainer.run(train_fn)
+    trainer.shutdown()
+
+    count = Counter(devices)
+    if num_gpus_per_worker == 0.5:
+        for i in range(4):
+            assert count[i] == 4
+    elif num_gpus_per_worker == 1:
+        for i in range(4):
+            assert count[i] == 2
+    elif num_gpus_per_worker == 2:
+        for i in range(2):
+            assert count[2 * i] == 2
+    else:
+        raise RuntimeError(
+            "New parameter for this test has been added without checking that the "
+            "correct devices have been returned."
+        )
+
+
 # TODO: Refactor as a backend test.
 def test_torch_prepare_model(ray_start_4_cpus_2_gpus):
     """Tests if ``prepare_model`` correctly wraps in DDP."""
@@ -346,6 +402,10 @@ def test_tune_fashion_mnist_gpu(ray_start_4_cpus_2_gpus):
     torch_fashion_mnist(num_workers=2, use_gpu=True, num_samples=1)
 
 
+def test_concurrent_tune_fashion_mnist_gpu(ray_start_4_cpus_2_gpus):
+    torch_fashion_mnist(num_workers=1, use_gpu=True, num_samples=2)
+
+
 def test_tune_tensorflow_mnist_gpu(ray_start_4_cpus_2_gpus):
     tune_tensorflow_mnist(num_workers=2, use_gpu=True, num_samples=1)
 

python/ray/train/tests/test_tune.py

@@ -3,6 +3,7 @@
 import pytest
 
 import ray
+from ray.cluster_utils import Cluster
 import ray.train as train
 from ray import tune
 from ray.tune import TuneError
@@ -23,6 +24,9 @@
 from ray.tune.tune_config import TuneConfig
 from ray.tune.tuner import Tuner
 
+from ray.train.torch import TorchConfig
+from unittest.mock import patch
+
 
 @pytest.fixture
 def ray_start_4_cpus():
@@ -40,6 +44,20 @@ def ray_start_8_cpus():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_2_node_4_gpu():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=2, num_gpus=4)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
 class TestConfig(BackendConfig):
     @property
     def backend_cls(self):
@@ -55,8 +73,6 @@ def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
 
 
 def torch_fashion_mnist(num_workers, use_gpu, num_samples):
-    epochs = 2
-
     trainer = TorchTrainer(
         fashion_mnist_train_func,
         scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
@@ -67,7 +83,7 @@ def torch_fashion_mnist(num_workers, use_gpu, num_samples):
             "train_loop_config": {
                 "lr": tune.loguniform(1e-4, 1e-1),
                 "batch_size": tune.choice([32, 64, 128]),
-                "epochs": epochs,
+                "epochs": 2,
             }
         },
         tune_config=TuneConfig(
@@ -86,8 +102,6 @@ def test_tune_torch_fashion_mnist(ray_start_8_cpus):
 
 
 def tune_tensorflow_mnist(num_workers, use_gpu, num_samples):
-    epochs = 2
-
     trainer = TensorflowTrainer(
         tensorflow_mnist_train_func,
         scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
@@ -98,7 +112,7 @@ def tune_tensorflow_mnist(num_workers, use_gpu, num_samples):
             "train_loop_config": {
                 "lr": tune.loguniform(1e-4, 1e-1),
                 "batch_size": tune.choice([32, 64, 128]),
-                "epochs": epochs,
+                "epochs": 2,
             }
         },
         tune_config=TuneConfig(
@@ -307,6 +321,47 @@ def train_func():
     assert len(trial_dfs[0]["training_iteration"]) == 4
 
 
+@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
+def test_tune_torch_get_device_gpu(ray_2_node_4_gpu, num_gpus_per_worker):
+    from ray import tune
+    from ray.tune.tuner import Tuner, TuneConfig
+
+    num_samples = 2
+
+    @patch("torch.cuda.is_available", lambda: True)
+    def train_func():
+        train.report(device_id=train.torch.get_device().index)
+
+    trainer = TorchTrainer(
+        train_func,
+        torch_config=TorchConfig(backend="gloo"),
+        scaling_config={
+            "num_workers": 2,
+            "use_gpu": True,
+            "resources_per_worker": {"GPU": num_gpus_per_worker},
+        },
+    )
+
+    tuner = Tuner(
+        trainer,
+        param_space={
+            "train_loop_config": {
+                "dummy": tune.choice([32, 64, 128]),
+            }
+        },
+        tune_config=TuneConfig(
+            num_samples=num_samples,
+        ),
+    )
+    analysis = tuner.fit()._experiment_analysis
+    trial_dfs = list(analysis.trial_dataframes.values())
+    device_ids = [trial_df["device_id"].tolist() for trial_df in trial_dfs]
+
+    assert len(device_ids) == num_samples
+    for i in range(num_samples):
+        assert device_ids[i][0] == 0
+
+
 if __name__ == "__main__":
     import sys
 

python/ray/train/torch/train_loop_utils.py

@@ -463,11 +463,48 @@ def wrapper(worker_id):
         return data_loader
 
     def get_device(self) -> torch.device:
-        """Gets the correct torch device to use for training."""
+        """Gets the correct torch device to use for training.
+
+        Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
+        superset of the `ray.get_gpu_ids()`.
+
+        Example:
+            >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "3,4"
+            >>> # ray.get_gpu_ids() == [3]
+            >>> # torch.cuda.is_available() == True
+            >>> # get_device() == torch.device("cuda:0")
+
+            >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4"
+            >>> # ray.get_gpu_ids() == [4]
+            >>> # torch.cuda.is_available() == True
+            >>> # get_device() == torch.device("cuda:4")
+
+            >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5"
+            >>> # ray.get_gpu_ids() == [4,5]
+            >>> # torch.cuda.is_available() == True
+            >>> # get_device() == torch.device("cuda:4")
+        """
         if torch.cuda.is_available():
+            # GPU IDs are assigned by Ray after you specify "use_gpu"
             gpu_ids = ray.get_gpu_ids()
+
             if len(gpu_ids) > 0:
-                device_id = gpu_ids[0]
+                # By default, there should only be one GPU ID if `use_gpu=True`.
+                # If there are multiple GPUs, use the first one.
+                # If using fractional GPUs, these IDs are not guaranteed
+                # to be unique across different processes.
+                gpu_id = gpu_ids[0]
+
+                cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+                if cuda_visible_str and cuda_visible_str != "NoDevFiles":
+                    cuda_visible_list = [
+                        int(dev) for dev in cuda_visible_str.split(",")
+                    ]
+                    device_id = cuda_visible_list.index(gpu_id)
+                else:
+                    raise RuntimeError(
+                        f"CUDA_VISIBLE_DEVICES set incorrectly: {cuda_visible_str}"
+                    )
             else:
                 # If called on the driver or outside of Ray Train, return the
                 # 0th device.