| import unittest | |
| from scripts.device import choose_runtime | |
| GIB = 1024 ** 3 | |
| class DevicePolicyTests(unittest.TestCase): | |
| def test_eight_gigabyte_cuda_uses_group_offload(self): | |
| policy = choose_runtime('cuda', 8 * GIB) | |
| self.assertEqual(policy['offload'], 'group') | |
| self.assertFalse(policy['use_stream']) | |
| def test_large_cuda_keeps_a_whole_component_on_gpu(self): | |
| policy = choose_runtime('cuda', 32 * GIB) | |
| self.assertEqual(policy['offload'], 'model') | |
| def test_apple_silicon_stays_on_the_eager_resident_path(self): | |
| policy = choose_runtime('mps', 24 * GIB) | |
| self.assertEqual(policy['offload'], 'resident') | |
| self.assertFalse(policy['use_stream']) | |
| self.assertEqual(policy['dtype'], 'bfloat16') | |
| def test_cpu_is_resident(self): | |
| self.assertEqual(choose_runtime('cpu', 64 * GIB)['offload'], 'resident') | |
| def test_unknown_device_is_rejected(self): | |
| with self.assertRaises(ValueError): | |
| choose_runtime('xpu', 8 * GIB) | |
| if __name__ == '__main__': | |
| unittest.main() | |