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)