| """One minted token runs a whole network of machines, each under its own name.""" |
| 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 OneTokenManyMachines(unittest.TestCase): |
| def test_mint_once_lease_from_many(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('{}') |
| 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({})) |
|
|
| def get(path, token=None): |
| req = urllib.request.Request(url + path, headers={'Authorization': f'Bearer {token}'} if token else {}) |
| try: |
| with opener.open(req, timeout=10) as r: |
| return r.status, r.read() |
| except urllib.error.HTTPError as e: |
| return e.code, e.read() |
|
|
| 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', '--join-bits', '4', '--min-steps', '1'], stdout=log, stderr=subprocess.STDOUT) |
| try: |
| deadline = time.monotonic() + 20 |
| while True: |
| try: |
| if get('/status')[0] == 200: |
| break |
| except urllib.error.URLError: |
| pass |
| 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) |
| path = root / 'worker.token' |
| r = subprocess.run([sys.executable, script, 'enrol', '--url', url, '--token-file', str(path)], |
| capture_output=True, text=True, timeout=60) |
| self.assertEqual(r.returncode, 0, r.stdout + r.stderr) |
| self.assertEqual(path.stat().st_mode & 0o777, 0o600) |
| token = path.read_text().strip() |
| self.assertNotIn(token, r.stdout + r.stderr) |
| |
| codes = [get(f'/lease?name=machine-{i}&session=s', token)[0] for i in range(6)] |
| self.assertEqual(codes, [200] * 6) |
| self.assertEqual(get('/lease?name=machine-0&session=s', token)[0], 429) |
| status = json.loads(get('/status')[1]) |
| self.assertEqual(status['open_join']['participants'], 1) |
| self.assertEqual(len(status['active_leases']), 6) |
| finally: |
| server.terminate() |
| try: |
| server.wait(timeout=10) |
| except subprocess.TimeoutExpired: |
| server.kill() |
| server.wait(timeout=5) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
