| """CPU regression tests and mocked backend selection; actual GPU test skips without a GPU. |
| |
| Run: python -m unittest -v |
| Mocks test routing only and are not evidence of AMD or NVIDIA hardware validation. |
| """ |
| import argparse |
| import contextlib |
| import io |
| import json |
| import math |
| import threading |
| import unittest |
| from unittest import mock |
|
|
| import torch |
| import dbmoe |
|
|
|
|
| def tiny_lease(): |
| cfg = dbmoe.make_cfg(dict(d=16, heads=2, layers_per_block=1, blocks=2, experts=2, ffn=32, ctx=8)) |
| return dict(cfg=cfg, block=0, block_version=0, sigmas=dbmoe.cfg_sigmas(cfg), lease='test-only', |
| shared=dbmoe.Shared(cfg).state_dict(), block_state=dbmoe.Block(cfg).state_dict(), |
| steps=2, batch=1, lr=1e-3, data=torch.arange(18, dtype=torch.uint8).reshape(2, 9)) |
|
|
|
|
| class DeviceSelectionTests(unittest.TestCase): |
| def setUp(self): |
| patcher = mock.patch.object(dbmoe, 'xpu_available', return_value=False) |
| patcher.start() |
| self.addCleanup(patcher.stop) |
|
|
| def fake_gpu(self, hip=None, available=True, count=2, current=0): |
| stack = contextlib.ExitStack() |
| stack.enter_context(mock.patch.object(torch.version, 'hip', hip)) |
| stack.enter_context(mock.patch.object(torch.cuda, 'is_available', return_value=available)) |
| stack.enter_context(mock.patch.object(torch.cuda, 'device_count', return_value=count)) |
| stack.enter_context(mock.patch.object(torch.cuda, 'current_device', return_value=current)) |
| self.addCleanup(stack.close) |
|
|
| def test_cpu_does_not_probe_cuda(self): |
| with mock.patch.object(torch.cuda, 'is_available', side_effect=AssertionError('GPU touched')): |
| self.assertEqual(dbmoe.resolve_worker_device('cpu'), |
| dict(torch_device='cpu', device='cpu', backend='cpu')) |
|
|
| def test_cpu_default_does_not_run_training_probe(self): |
| with mock.patch.object(dbmoe, 'device_smoke_test') as probe: |
| dbmoe.prepare_worker_device('cpu') |
| probe.assert_not_called() |
|
|
| def test_auto_cpu(self): |
| self.fake_gpu(available=False) |
| with mock.patch.object(dbmoe, 'mps_available', return_value=False): |
| self.assertEqual(dbmoe.resolve_worker_device('auto')['backend'], 'cpu') |
|
|
| def test_nvidia_cuda_and_index(self): |
| self.fake_gpu() |
| self.assertEqual(dbmoe.resolve_worker_device('cuda:1'), |
| dict(torch_device='cuda:1', device='cuda:1', backend='cuda')) |
| self.assertEqual(dbmoe.resolve_worker_device('cuda')['torch_device'], 'cuda:0') |
| self.assertEqual(dbmoe.resolve_worker_device('auto')['backend'], 'cuda') |
|
|
| def test_rocm_aliases_and_legacy_cuda(self): |
| self.fake_gpu(hip='test-hip') |
| for spec in ('rocm:1', 'hip:1', 'cuda:1'): |
| with self.subTest(spec=spec): |
| self.assertEqual(dbmoe.resolve_worker_device(spec), |
| dict(torch_device='cuda:1', device='rocm:1', backend='rocm')) |
| for spec in ('rocm', 'hip', 'cuda', 'auto'): |
| self.assertEqual(dbmoe.resolve_worker_device(spec)['backend'], 'rocm') |
|
|
| def test_unavailable_explicit_gpu_never_falls_back(self): |
| self.fake_gpu(hip='test-hip', available=False) |
| for spec in ('cuda', 'cuda:0', 'rocm', 'rocm:0', 'hip'): |
| with self.subTest(spec=spec), self.assertRaisesRegex(RuntimeError, 'cannot access a GPU'): |
| dbmoe.resolve_worker_device(spec) |
|
|
| def test_rocm_rejects_non_hip_build_even_with_nvidia_gpu(self): |
| self.fake_gpu() |
| with self.assertRaisesRegex(RuntimeError, 'no HIP support'): |
| dbmoe.resolve_worker_device('rocm') |
|
|
| def test_invalid_and_out_of_range_devices(self): |
| self.fake_gpu(hip='test-hip') |
| for spec in ('', 'cpu:0', 'cuda:', 'rocm:-1', 'cuda:abc', 'rocm:1:2', 'cuda:²', 'metal', 'auto:1'): |
| with self.subTest(spec=spec), self.assertRaises(ValueError): |
| dbmoe.resolve_worker_device(spec) |
| with self.assertRaisesRegex(ValueError, 'out of range'): |
| dbmoe.resolve_worker_device('rocm:2') |
|
|
| def test_selected_device_is_used_for_gpu_probe(self): |
| self.fake_gpu(hip='test-hip') |
| with mock.patch.object(torch.cuda, 'set_device') as select, \ |
| mock.patch.object(torch.cuda, 'get_device_name', return_value='Mock AMD GPU'), \ |
| mock.patch.object(dbmoe, 'device_smoke_test', return_value=[1., 1.]) as probe: |
| info = dbmoe.prepare_worker_device('rocm:1') |
| select.assert_called_once_with('cuda:1') |
| probe.assert_called_once_with('cuda:1') |
| self.assertEqual(info['device'], 'rocm:1') |
|
|
| def test_probe_failure_is_fatal_before_any_lease_or_token_read(self): |
| a = argparse.Namespace(threads=1, device='rocm', token='@missing-token-file') |
| with mock.patch.object(dbmoe, 'prepare_worker_device', side_effect=RuntimeError('bad driver')), \ |
| mock.patch.object(dbmoe, 'http') as http: |
| with self.assertRaisesRegex(SystemExit, 'bad driver'): |
| dbmoe.worker(a) |
| http.assert_not_called() |
|
|
|
|
| class TrainingTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| torch.set_num_threads(1) |
|
|
| def test_offline_cpu_probe_without_network(self): |
| with mock.patch.object(dbmoe, 'http', side_effect=AssertionError('network touched')): |
| losses = dbmoe.device_smoke_test('cpu') |
| self.assertEqual(len(losses), 2) |
| self.assertTrue(all(math.isfinite(v) for v in losses)) |
|
|
| def test_sub_centisecond_lease_reports_positive_duration(self): |
| L = tiny_lease() |
| with mock.patch.object(dbmoe.time, 'perf_counter', side_effect=[100., 100.001]): |
| up = dbmoe.from_bytes(dbmoe.to_bytes(dbmoe.train_lease(L, 'cpu'))) |
| self.assertEqual(up['seconds'], .01) |
| C = dbmoe.Coordinator.__new__(dbmoe.Coordinator) |
| C.lr, C.model = L['lr'], dbmoe.Model(L['cfg']) |
| self.assertIsNone(C._check(up, L)) |
|
|
| def test_cpu_train_delta_format_cache_and_coordinator_validation(self): |
| torch.manual_seed(123) |
| L = tiny_lease() |
| original = {section: {k: v.clone() for k, v in L[section].items()} |
| for section in ('shared', 'block_state')} |
| cache = {} |
| up = dbmoe.from_bytes(dbmoe.to_bytes(dbmoe.train_lease(L, 'cpu', cache))) |
| self.assertEqual(up['tokens'], 16) |
| self.assertEqual(len(up['losses']), 2) |
| self.assertTrue(up['fresh_optimizer']) |
| C = dbmoe.Coordinator.__new__(dbmoe.Coordinator) |
| C.lr, C.model = L['lr'], dbmoe.Model(L['cfg']) |
| self.assertIsNone(C._check(up, L)) |
| for section in original: |
| for key in original[section]: |
| torch.testing.assert_close(L[section][key], original[section][key], rtol=0, atol=0) |
| self.assertTrue(any(bool(t.abs().max() > 0) for t in up['d_block'].values())) |
| for section in ('d_block', 'd_shared'): |
| for t in up[section].values(): |
| self.assertEqual(t.device.type, 'cpu') |
| self.assertEqual(t.dtype, torch.float32) |
| self.assertTrue(bool(torch.isfinite(t).all())) |
| L['block_version'] = 1 |
| second = dbmoe.train_lease(L, 'cpu', cache) |
| self.assertFalse(second['fresh_optimizer']) |
| self.assertIsNone(C._check(second, L)) |
|
|
| def test_worker_one_cpu_lease_round_trip(self): |
| L = tiny_lease() |
| calls = [] |
| def http(method, url, token, body=None): |
| calls.append((method, url, token, body)) |
| return dbmoe.to_bytes(L) if method == 'GET' else json.dumps(dict(ok=True)).encode() |
| a = argparse.Namespace(threads=1, device='cpu', url='https://example.invalid', token='test-only', |
| name='test', hours=0, steps=2, batch=1, once=1) |
| with mock.patch.object(dbmoe, 'http', side_effect=http), contextlib.redirect_stdout(io.StringIO()): |
| dbmoe.worker(a) |
| self.assertEqual([c[0] for c in calls], ['GET', 'POST']) |
| self.assertEqual(dbmoe.from_bytes(calls[1][3])['device'], 'cpu') |
|
|
| def test_amd_payload_uses_rocm_label_not_pytorch_cuda_label(self): |
| L = tiny_lease() |
| up = dbmoe.train_lease(L, 'cpu') |
| sent = [] |
| def http(method, url, token, body=None): |
| if body: |
| sent.append(dbmoe.from_bytes(body)) |
| return dbmoe.to_bytes(L) if method == 'GET' else b'{"ok":true}' |
| a = argparse.Namespace(threads=1, device='rocm', url='https://example.invalid', token='test-only', |
| name='test', hours=0, steps=2, batch=1, once=1) |
| info = dict(torch_device='cuda:0', device='rocm:0', backend='rocm', gpu_name='Mock GPU') |
| with mock.patch.object(dbmoe, 'prepare_worker_device', return_value=info), \ |
| mock.patch.object(dbmoe, 'train_lease', return_value=up) as train, \ |
| mock.patch.object(dbmoe, 'http', side_effect=http), contextlib.redirect_stdout(io.StringIO()): |
| dbmoe.worker(a) |
| self.assertEqual(train.call_args.args[1], 'cuda:0') |
| self.assertEqual(sent[0]['device'], 'rocm:0') |
|
|
| @unittest.skipUnless(torch.cuda.is_available(), 'No CUDA/ROCm GPU available; hardware test not run') |
| def test_real_gpu_forward_backward_and_serialization(self): |
| kind = 'rocm' if torch.version.hip else 'cuda' |
| info = dbmoe.prepare_worker_device(kind) |
| self.assertEqual(info['backend'], kind) |
| self.assertEqual(len(info['check_losses']), 2) |
|
|
|
|
| class HardwarePrivacyTests(unittest.TestCase): |
| def test_cli_default_name_does_not_read_hostname(self): |
| with mock.patch.object(dbmoe.sys, 'argv', ['dbmoe.py', 'worker', '--url', 'https://example.invalid', |
| '--token', 'test-only']), \ |
| mock.patch.object(dbmoe.socket, 'gethostname', side_effect=AssertionError('hostname touched')), \ |
| mock.patch.object(dbmoe, 'worker') as worker: |
| dbmoe.main() |
| self.assertTrue(worker.call_args.args[0].name.startswith('worker-')) |
|
|
| def test_status_filters_legacy_host_and_ip(self): |
| C = dbmoe.Coordinator.__new__(dbmoe.Coordinator) |
| C.lock = threading.RLock() |
| C.model = dbmoe.Model(tiny_lease()['cfg']) |
| C.meta = dict(updates=0, tokens=0, rejected=0, block_updates=[0, 0], block_ce=[None, None], |
| evals=[], workers={'pseudonym': dict(host='private-host', ip='192.0.2.1', device='cpu')}) |
| C.leases = {} |
| C.data = argparse.Namespace(docs=0, buf=b'', err=None) |
| C.steps, C.batch, C.lr, C.max_steps, C.max_batch, C.ttl = 2, 1, .001, 10, 2, 60 |
| self.assertEqual(C.status()['workers']['pseudonym'], dict(device='cpu')) |
| self.assertIn('host', C.meta['workers']['pseudonym']) |
|
|
| def test_allowlist_discards_private_and_unknown_fields(self): |
| report = dbmoe.normalize_hardware(dict(schema=1, self_reported=True, backend='rocm', |
| cpu_arch='x86_64', cpu_threads=2, torch_version='test', runtime_version='test', |
| gpu_name='AMD\n Test', gpu_memory_gib=15.96, host='secret-host', ip='192.0.2.1', |
| serial='secret-serial', token='private-token', path='/private/path')) |
| self.assertEqual(set(report), {'schema', 'self_reported', 'backend', 'cpu_arch', 'cpu_threads', |
| 'torch_version', 'runtime_version', 'gpu_name', 'gpu_memory_gib'}) |
| self.assertEqual(report['gpu_name'], 'AMD Test') |
| self.assertEqual(report['gpu_memory_gib'], 16.0) |
|
|
| def test_malformed_metadata_is_ignored_without_raising(self): |
| for report in (None, [], 'text', {'schema': torch.tensor([1, 2])}, |
| {'schema': True}, {'schema': 1, 'self_reported': False}, |
| {'schema': 1, 'self_reported': True, 'backend': torch.tensor([1, 2])}): |
| with self.subTest(report=report): |
| self.assertIsNone(dbmoe.normalize_hardware(report)) |
| for memory in (float('nan'), float('inf'), -1, True, [16], torch.tensor([16, 16])): |
| report = dbmoe.normalize_hardware(dict(schema=1, self_reported=True, backend='rocm', |
| gpu_memory_gib=memory, cpu_threads=True)) |
| self.assertNotIn('gpu_memory_gib', report) |
| self.assertNotIn('cpu_threads', report) |
|
|
| def test_cpu_hardware_never_queries_gpu_or_hostname(self): |
| with mock.patch.object(torch.cuda, 'get_device_properties', side_effect=AssertionError('GPU touched')), \ |
| mock.patch.object(dbmoe.socket, 'gethostname', side_effect=AssertionError('hostname touched')): |
| report = dbmoe.worker_hardware(dict(backend='cpu', torch_device='cpu'), 2) |
| self.assertEqual(report['cpu_threads'], 2) |
| self.assertEqual(report['backend'], 'cpu') |
| self.assertNotIn('gpu_name', report) |
|
|
| def test_gpu_hardware_queries_selected_device_only(self): |
| props = argparse.Namespace(total_memory=24 * 2**30) |
| with mock.patch.object(torch.cuda, 'get_device_properties', return_value=props) as get_props: |
| report = dbmoe.worker_hardware(dict(backend='rocm', torch_device='cuda:1', gpu_name='Mock AMD'), 3) |
| get_props.assert_called_once_with('cuda:1') |
| self.assertEqual(report['gpu_memory_gib'], 24.) |
| self.assertEqual(report['gpu_name'], 'Mock AMD') |
|
|
| def test_worker_metadata_requires_opt_in_and_never_sends_host(self): |
| L = tiny_lease() |
| up = dbmoe.train_lease(L, 'cpu') |
| for opted_in in (False, True): |
| calls = [] |
| def http(method, url, token, body=None): |
| if body: |
| calls.append(dbmoe.from_bytes(body)) |
| return dbmoe.to_bytes(L) if method == 'GET' else b'{"ok":true}' |
| a = argparse.Namespace(threads=1, device='cpu', url='https://example.invalid', token='test-only', |
| name='test', hours=0, steps=2, batch=1, once=1, no_share_hardware=not opted_in) |
| with self.subTest(opted_in=opted_in), \ |
| mock.patch.object(dbmoe, 'train_lease', return_value=dict(up)), \ |
| mock.patch.object(dbmoe, 'http', side_effect=http), contextlib.redirect_stdout(io.StringIO()): |
| dbmoe.worker(a) |
| self.assertNotIn('host', calls[0]) |
| self.assertNotIn('ip', calls[0]) |
| self.assertEqual('hardware' in calls[0], opted_in) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
|
|
|
|
