| """Isolated Chrome WebGPU service and OpenAI bridge; no cloud/model fallback.""" |
| import asyncio |
| import argparse |
| import importlib.util |
| import json |
| import mimetypes |
| from pathlib import Path |
| import shutil |
| import threading |
| import time |
| import uuid |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
|
|
| from openai_server import normalize_messages, parse_tools |
| from playwright.async_api import async_playwright |
|
|
| ROOT = Path(r'D:\webbrain-training\compass-tiny-v3-spark-x25-20260926\onnx-v3') |
| SOURCE = Path(__file__).parent |
| MODEL_ID = 'webbrain-compass-tiny-v3-onnx-fp16' |
| PORT = 8561 |
| VENDOR = Path(r'F:\webbrain-data\compass-tiny-v2-webgpu-recovery-v24\runtime\vendor') |
|
|
|
|
| async def main(): |
| global ROOT, PORT |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--run-tag', default='browser-v5') |
| parser.add_argument('--artifact-root', type=Path, default=ROOT) |
| parser.add_argument('--package-mode', action='store_true', help='Root is the published package, not the export work directory') |
| parser.add_argument('--port', type=int, default=PORT) |
| args = parser.parse_args() |
| ROOT, PORT = args.artifact_root.resolve(), args.port |
| model_dir = ROOT if args.package_mode else ROOT / 'bundle' |
| run = ROOT / args.run_tag |
| run.mkdir(exist_ok=False) |
| runtime = ROOT / 'runtime' |
| runtime.mkdir(exist_ok=True) |
| if not (runtime / 'vendor').exists(): |
| if not VENDOR.is_dir(): |
| raise RuntimeError('Missing pinned runtime/vendor files; download the complete package') |
| shutil.copytree(VENDOR, runtime / 'vendor') |
| dxgi_source = SOURCE / 'compass_dxgi_5090.py' |
| if not dxgi_source.is_file(): |
| dxgi_source = Path(r'F:\webbrain-data\webbrain-data\scripts\hf\compass_dxgi_5090.py') |
| module_spec = importlib.util.spec_from_file_location('dxgi', dxgi_source) |
| dxgi = importlib.util.module_from_spec(module_spec) |
| module_spec.loader.exec_module(dxgi) |
| adapter = dxgi.discover() |
| loop = asyncio.get_running_loop() |
| lock = asyncio.Lock() |
| state = {'ready': False, 'model': MODEL_ID, 'precision': 'FP16 ONNX; FP32 residual/norm/RoPE/softmax', |
| 'runtime': 'ONNX Runtime WebGPU', 'gpu': 'RTX 5090', 'cloudFallback': False, |
| 'maxContextTokens': 4096, 'temperature': 'greedy-only benchmark endpoint'} |
| page = None |
|
|
| async def rpc(method, payload): |
| async with lock: |
| return await page.evaluate('async ({method,payload})=>window[method](payload)', {'method': method, 'payload': payload}) |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def log_message(self, fmt, *args): |
| if self.path.startswith(('/v1/', '/infer', '/template', '/health')): |
| print(json.dumps({'http': fmt % args, 'time': time.time()}), flush=True) |
|
|
| def json(self, status, value): |
| data = json.dumps(value, ensure_ascii=False).encode() |
| self.send_response(status) |
| self.send_header('Content-Type', 'application/json') |
| self.send_header('Content-Length', str(len(data))) |
| self.end_headers() |
| self.wfile.write(data) |
|
|
| def do_GET(self): |
| if self.path == '/health': |
| return self.json(200, state) |
| if self.path == '/v1/models': |
| return self.json(200, {'object': 'list', 'data': [{'id': MODEL_ID, 'object': 'model', 'owned_by': 'webbrain-one', **state}]}) |
| if self.path == '/abi': |
| return self.json(200, json.loads((ROOT / 'graph-abi.json').read_text())) |
| if self.path == '/': |
| data = b'<html><script type="module" src="/browser_runtime.js"></script></html>' |
| self.send_response(200) |
| self.send_header('Content-Type', 'text/html') |
| self.send_header('Content-Length', str(len(data))) |
| self.end_headers() |
| return self.wfile.write(data) |
| if self.path == '/browser_runtime.js': |
| target = SOURCE / 'browser_runtime.js' |
| elif self.path.startswith('/model/'): |
| target = model_dir / self.path[len('/model/'):] |
| elif self.path.startswith('/runtime/'): |
| target = ROOT / self.path.lstrip('/') |
| else: |
| return self.json(404, {'error': 'Not found'}) |
| target = target.resolve() |
| allowed = [ROOT.resolve(), SOURCE.resolve()] |
| if not any(target.is_relative_to(p) for p in allowed) or not target.is_file(): |
| return self.json(404, {'error': 'Not found'}) |
| self.send_response(200) |
| suffix = target.suffix |
| self.send_header('Content-Type', {'mjs': 'text/javascript', 'js': 'text/javascript', 'wasm': 'application/wasm'}.get(suffix[1:], mimetypes.guess_type(str(target))[0] or 'application/octet-stream')) |
| self.send_header('Content-Length', str(target.stat().st_size)) |
| self.end_headers() |
| with target.open('rb') as f: |
| shutil.copyfileobj(f, self.wfile, length=4 * 1024 * 1024) |
|
|
| def do_POST(self): |
| try: |
| length = int(self.headers.get('Content-Length', '0')) |
| if not 0 < length <= 4 * 1024 * 1024: |
| raise ValueError('Invalid body size') |
| body = json.loads(self.rfile.read(length)) |
| if not state['ready']: |
| return self.json(503, {'error': 'WebGPU model not ready'}) |
| if self.path == '/infer': |
| method, payload = 'sparkNumerics', body |
| elif self.path == '/template': |
| method, payload = 'sparkTemplate', body |
| elif self.path == '/v1/chat/completions': |
| if body.get('model') != MODEL_ID or body.get('stream', False) or float(body.get('temperature', 0)) != 0: |
| raise ValueError('Use exact ONNX model id, stream=false, temperature=0') |
| maximum = int(body.get('max_completion_tokens', body.get('max_tokens', 512))) |
| if not 1 <= maximum <= 2048: |
| raise ValueError('max_tokens must be 1..2048') |
| method, payload = 'sparkGenerate', {'messages': normalize_messages(body['messages']), |
| 'tools': body.get('tools'), 'maxTokens': maximum} |
| else: |
| return self.json(404, {'error': 'Not found'}) |
| result = asyncio.run_coroutine_threadsafe(rpc(method, payload), loop).result(timeout=600) |
| if self.path == '/v1/chat/completions': |
| raw = result['content'] |
| parser_error = None |
| try: |
| calls, content = parse_tools(raw, body.get('tools')) |
| except Exception as exc: |
| calls, content, parser_error = [], raw, str(exc) |
| message = {'role': 'assistant', 'content': content or None} |
| finish = result['finishReason'] |
| if calls: |
| message['tool_calls'] = calls |
| finish = 'tool_calls' |
| result = {'id': 'chatcmpl-' + uuid.uuid4().hex, 'object': 'chat.completion', 'created': int(time.time()), |
| 'model': MODEL_ID, 'choices': [{'index': 0, 'message': message, 'finish_reason': finish}], |
| 'usage': {'prompt_tokens': result['promptTokens'], 'completion_tokens': result['completionTokens'], |
| 'total_tokens': result['promptTokens'] + result['completionTokens']}, 'x_runtime': state} |
| if parser_error: |
| result['x_parser_error'] = parser_error |
| self.json(200, result) |
| except (ValueError, KeyError, TypeError) as exc: |
| self.json(400, {'error': {'message': str(exc), 'type': 'invalid_request_error'}}) |
| except Exception as exc: |
| self.json(500, {'error': {'message': str(exc), 'type': 'runtime_error'}}) |
|
|
| server = ThreadingHTTPServer(('127.0.0.1', PORT), Handler) |
| threading.Thread(target=server.serve_forever, daemon=True).start() |
| async with async_playwright() as p: |
| browser = await p.chromium.launch(executable_path=r'C:\Program Files\Google\Chrome\Application\chrome.exe', headless=True, |
| args=['--enable-unsafe-webgpu', '--ignore-gpu-blocklist', '--use-angle=d3d11', |
| f'--use-adapter-luid={adapter["high"]},{adapter["low"]}', '--force_high_performance_gpu', |
| '--enable-features=EGLDualGPURendering', '--enable-dawn-features=allow_unsafe_apis', |
| '--js-flags=--max-old-space-size=16384']) |
| try: |
| cdp = await browser.new_browser_cdp_session() |
| info = await cdp.send('SystemInfo.getInfo') |
| renderer = info['gpu']['auxAttributes']['glRenderer'] |
| if 'RTX 5090' not in renderer: |
| raise RuntimeError('Chrome selected the wrong GPU') |
| (run / 'gpu-identity.json').write_text(json.dumps({'dxgi': adapter, 'renderer': renderer}), encoding='utf-8') |
| page = await browser.new_page() |
| page.on('requestfailed', lambda request: print(json.dumps({'requestFailed': request.url.split('?')[0], 'reason': request.failure}), flush=True)) |
| page.on('console', lambda message: print(json.dumps({'console': message.type, 'text': message.text[:1500]}), flush=True)) |
| page.on('pageerror', lambda error: print(json.dumps({'pageError': str(error)}), flush=True)) |
| await page.goto(f'http://127.0.0.1:{PORT}/') |
| await page.wait_for_function('typeof window.initializeSpark === "function"') |
| ready = await page.evaluate('window.initializeSpark()') |
| state.update(ready=True, browser=browser.version, **ready) |
| (run / 'ready.json').write_text(json.dumps(state), encoding='utf-8') |
| print(json.dumps({'ready': True, 'baseUrl': f'http://127.0.0.1:{PORT}/v1', **state}), flush=True) |
| await asyncio.Event().wait() |
| finally: |
| server.shutdown() |
| await browser.close() |
|
|
|
|
| if __name__ == '__main__': |
| asyncio.run(main()) |
|
|
