File size: 1,796 Bytes
4a7414e | 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 36 37 38 39 | """End to end: the Rust worker (worker-rs) trains leases from a real local coordinator over the safetensors wire format.
Skipped unless DBMOE_RUST_WORKER names a built dbmoe-worker binary (the worker-rs CI job sets it):
cargo build --release --manifest-path worker-rs/Cargo.toml
DBMOE_RUST_WORKER=worker-rs/target/release/dbmoe-worker python -m unittest test_rust_worker_loopback -v
"""
import json
import os
import subprocess
import unittest
from test_safetensors_wire import start_coordinator
class RustWorkerLoopbackTest(unittest.TestCase):
def test_rust_worker_completes_leases(self):
exe = os.environ.get('DBMOE_RUST_WORKER')
if not exe:
self.skipTest('set DBMOE_RUST_WORKER to a built dbmoe-worker binary to run this test')
url, opener, root = start_coordinator(self)
home = str(root / 'home') # keep ~/.dbmoe out of the real home directory
env = dict(os.environ, HOME=home, USERPROFILE=home)
cmd = [exe, 'worker', '--url', url, '--token', 'local-test-only', '--name', 'rust-test', '--device', 'cpu',
'--steps', '2', '--batch', '1', '--once', '2', '--threads', '1', '--seed', '1']
r = subprocess.run(cmd, capture_output=True, text=True, timeout=600, env=env)
self.assertEqual(r.returncode, 0, r.stdout + r.stderr)
with opener.open(url + '/status', timeout=5) as resp:
status = json.load(resp)
self.assertEqual(status['updates'], 2, r.stdout + r.stderr)
self.assertEqual(r.stdout.count('"ok":true'), 2, r.stdout)
worker = next(iter(status['workers'].values()))
self.assertEqual(worker['device'], 'cpu')
self.assertNotIn('hardware', worker)
if __name__ == '__main__':
unittest.main()
|