def test1VisibleNotInitialized(self):
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
with patch("torch.cuda.is_initialized") as init_mock:
init_mock.return_value = False
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 test2VisibleNotInitialized(self):
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
with patch("torch.cuda.is_initialized") as init_mock:
init_mock.return_value = False
mock_runner = MagicMock()
mock_runner._set_cuda_device = MagicMock()
LocalDistributedRunner._try_reserve_and_set_cuda(mock_runner)
args, _ = mock_runner._set_cuda_device.call_args
self.assertTrue(("1" in args) or "0" in args)
self.assertEquals(len(os.environ["CUDA_VISIBLE_DEVICES"]), 1)
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 _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)
