| """End-to-end test against an isolated, tiny coordinator bound only to loopback.""" |
| import json |
| from pathlib import Path |
| import socket |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| import unittest |
| import urllib.error |
| import urllib.request |
|
|
|
|
| class LoopbackIntegrationTest(unittest.TestCase): |
| def test_cpu_leases_opt_in_and_withdrawal(self): |
| script = str(Path(__file__).with_name('dbmoe.py').resolve()) |
| with tempfile.TemporaryDirectory(prefix='dbmoe-test-') as tmp: |
| root = Path(tmp) |
| (root / 'training.txt').write_text('This is local synthetic training data only.\n' * 100) |
| (root / 'tokens.json').write_text(json.dumps({'local-test-only': 'local'})) |
| (root / 'worker.token').write_text('local-test-only') |
| cfg = dict(d=16, heads=2, layers_per_block=1, blocks=2, experts=2, ffn=32, ctx=8) |
| with socket.socket() as s: |
| s.bind(('127.0.0.1', 0)) |
| port = s.getsockname()[1] |
| url = f'http://127.0.0.1:{port}' |
| opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) |
| with (root / 'server.log').open('w+') as log: |
| server = subprocess.Popen([sys.executable, script, 'coord', '--dir', str(root / 'run'), |
| '--host', '127.0.0.1', '--port', str(port), '--tokens', str(root / 'tokens.json'), |
| '--text-file', str(root / 'training.txt'), '--cfg', json.dumps(cfg), '--steps', '2', |
| '--batch', '1', '--heldout', '2', '--eval-every', '0', '--snap-every', '0', |
| '--eval-threads', '1', '--no-gate-all', '--min-steps', '1'], stdout=log, stderr=subprocess.STDOUT) |
| try: |
| deadline = time.monotonic() + 20 |
| while True: |
| try: |
| with opener.open(url + '/status', timeout=1) as response: |
| json.load(response) |
| break |
| except (urllib.error.URLError, TimeoutError): |
| if server.poll() is not None or time.monotonic() > deadline: |
| log.seek(0) |
| self.fail('Local coordinator did not start: ' + log.read()) |
| time.sleep(.1) |
| for opt_in in (True, False): |
| cmd = [sys.executable, script, 'worker', '--url', url, |
| '--token', '@' + str(root / 'worker.token'), '--name', 'public-test', |
| '--device', 'cpu', '--threads', '1', '--steps', '2', '--batch', '1', '--once', '1', '--no-serve'] |
| if not opt_in: |
| cmd.append('--no-share-hardware') |
| worker = subprocess.run(cmd, capture_output=True, text=True, timeout=20) |
| self.assertEqual(worker.returncode, 0, worker.stdout + worker.stderr) |
| with opener.open(url + '/status', timeout=2) as response: |
| status = json.load(response) |
| self.assertEqual(len(status['workers']), 1, status['workers']) |
| stats = next(iter(status['workers'].values())) |
| self.assertEqual('hardware' in stats, opt_in) |
| self.assertNotIn('host', stats) |
| self.assertNotIn('ip', stats) |
| if opt_in: |
| self.assertTrue(stats['hardware']['self_reported']) |
| self.assertEqual(stats['hardware']['backend'], 'cpu') |
| self.assertEqual(status['updates'], 2) |
| self.assertEqual(status['rejected'], 0) |
| self.assertEqual(status['tokens'], 32) |
| finally: |
| server.terminate() |
| try: |
| server.wait(timeout=10) |
| except subprocess.TimeoutExpired: |
| server.kill() |
| server.wait(timeout=5) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
