def testV2ReserveCPUResources(self):
mock_runner = MagicMock()
mock_runner._set_cpu_devices = MagicMock()
# reserve CPU and GPU
LocalDistributedRunner._try_reserve_and_set_resources(
mock_runner, 4, 1)
remaining = ray.available_resources()["CPU"]
self.assertEqual(int(remaining), 6)
def _testReserveCUDAResource(self, init_mock, num_cpus):
mock_runner = MagicMock()
mock_runner._set_cuda_device = MagicMock()
preset_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
LocalDistributedRunner._try_reserve_and_set_resources(
mock_runner, num_cpus, 1)
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.assertEqual(len(env_set_device), 1)
if preset_devices:
visible_devices = preset_devices.split(",")
self.assertIn(env_set_device, visible_devices)
device_int = int(local_device)
self.assertLess(device_int, len(visible_devices))
self.assertEqual(len(os.environ["CUDA_VISIBLE_DEVICES"]), 1)
else:
self.assertEqual(local_device, env_set_device)
