[sgd] Fix GPU Reservations in LocalDistributedRunner (#8157)

python/ray/util/sgd/tests/test_torch_runner.py

@@ -1,10 +1,14 @@
 import numpy as np
+import os
 import torch
 import torch.nn as nn
 import unittest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
+import ray
 from ray.util.sgd.torch.training_operator import TrainingOperator
+from ray.util.sgd.torch.distributed_torch_runner import (
+    LocalDistributedRunner, clear_dummy_actor)
 from ray.util.sgd.torch.torch_runner import TorchRunner
 
 
@@ -170,3 +174,81 @@ class TestTorchRunner(unittest.TestCase):
 
         with self.assertRaises(ValueError):
             runner.setup()
+
+
+class TestLocalDistributedRunner(unittest.TestCase):
+    def setUp(self):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        ray.init(num_gpus=4)
+
+    def tearDown(self):
+        clear_dummy_actor()
+        ray.shutdown()
+
+    def _testWithInitialized(self, init_mock):
+        mock_runner = MagicMock()
+        mock_runner._set_cuda_device = MagicMock()
+        preset_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+
+        LocalDistributedRunner._try_reserve_and_set_cuda(mock_runner)
+
+        self.assertTrue(mock_runner._set_cuda_device.called)
+        local_device = mock_runner._set_cuda_device.call_args[0][0]
+        env_set_device = os.environ["CUDA_VISIBLE_DEVICES"]
+        self.assertEquals(len(env_set_device), 1)
+
+        if preset_devices:
+            self.assertIn(env_set_device, preset_devices.split(","))
+            self.assertEquals(local_device, "0")
+        else:
+            self.assertEquals(local_device, env_set_device)
+
+    def testNoVisibleWithInitialized(self):
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = True
+            self._testWithInitialized(init_mock)
+
+    def test2VisibleWithInitialized(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = True
+            self._testWithInitialized(init_mock)
+
+    def test1VisibleWithInitialized(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = True
+            self._testWithInitialized(init_mock)
+
+    def _testNotInitialized(self, init_mock):
+        mock_runner = MagicMock()
+        mock_runner._set_cuda_device = MagicMock()
+        LocalDistributedRunner._try_reserve_and_set_cuda(mock_runner)
+        mock_runner._set_cuda_device.assert_called_with("0")
+        self.assertEquals(len(os.environ["CUDA_VISIBLE_DEVICES"]), 1)
+
+    def testNoVisibleNotInitialized(self):
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = False
+            self._testNotInitialized(init_mock)
+
+    def test2VisibleNotInitialized(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = False
+            self._testNotInitialized(init_mock)
+
+    def test1VisibleNotInitialized(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+        with patch("torch.cuda.is_initialized") as init_mock:
+            init_mock.return_value = False
+            self._testNotInitialized(init_mock)
+
+    @patch("torch.cuda.set_device")
+    def testSetDevice(self, set_mock):
+        mock_runner = MagicMock()
+        mock_runner._is_set = False
+        LocalDistributedRunner._set_cuda_device(mock_runner, "123")
+        self.assertEquals(mock_runner.local_device, "123")
+        self.assertTrue(set_mock.called)
+        set_mock.assert_called_with(123)