File size: 10,269 Bytes
c5cfd97 | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """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())
|