File size: 1,073 Bytes
9116984 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | 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()
|