File size: 9,605 Bytes
952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 952e1e7 7027070 | 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 | # SPDX-License-Identifier: MIT
"""Decision Index adapter. Start one persistent isolated GPU backend per engine."""
import atexit
import json
import os
from pathlib import Path
import select
import subprocess
import sys
import threading
import time
from decision_index.engines import Engine, Unsupported
ROOT = Path(__file__).resolve().parent
PROXY_MODEL = 'jev-1.13.0'
class EngineUnavailable(RuntimeError):
"""The native backend is gone (timeout, GPU failure or crash); this engine cannot serve until restarted."""
def proxy_tokens(state, questions):
"""A character-count size proxy: length of the compact request JSON // 4 (the literal model name is fixed, not looked up). No tokenizer, no clock, no answer-dependent step."""
return max(1, len(json.dumps({'model': PROXY_MODEL, 'state': state, 'questions': questions},
ensure_ascii=False, separators=(',', ':'))) // 4)
class Reax13Engine(Engine):
name = 'reax13'
latency = ('Synchronous request wall time: JSON IPC, both native prompt renderers, tokenization, '
'two FP8 forwards on GPU0, CUDA synchronization and equal probability averaging; '
'excludes loading and warmup. One external request at a time.')
def __init__(self, config=None, model=None, **options):
super().__init__(**options)
self.config = json.loads(Path(config or ROOT / 'reax13.json').read_text())
self.config.update(options)
self.max_input_proxy_tokens = int(self.config['max_input_proxy_tokens'])
if self.max_input_proxy_tokens <= 0:
raise ValueError('max_input_proxy_tokens must be positive')
if self.config['mode'] != 'reax13' or self.config['quantization'] != 'fp8':
raise ValueError('this engine supports reax13 FP8 only')
if self.config['optimization'] not in ('fastconv_longcache', 'plain'):
raise ValueError('unmeasured optimization requested')
if os.environ.get('DUO_ISOLATED') != '1' or not Path('/.dockerenv').exists():
raise RuntimeError('run inside the supplied isolated GPU container')
model_root = Path(model or self.config['model_root']).resolve()
paths = {part: (model_root / self.config[part]).resolve() for part in ('part_g', 'part_q')}
if any(not p.is_dir() for p in paths.values()):
raise FileNotFoundError('model root must contain both part-g/ and part-q/')
env = dict(os.environ, REAX_PART_G=str(paths['part_g']), REAX_PART_Q=str(paths['part_q']))
self.lock = threading.Lock()
self._buf = b''
self.process = subprocess.Popen([sys.executable, '-u', str(ROOT / 'backend.py')],
stdin=subprocess.PIPE, stdout=subprocess.PIPE,
text=True, bufsize=1, env=env)
atexit.register(self.close)
try:
ready = self._read(float(self.config['startup_timeout']))
if not ready.get('ready'):
raise RuntimeError('backend did not become ready: ' + str(ready))
self.ready = ready
except BaseException:
self.close()
raise
self.provenance = {'mode': 'reax13', 'config': self.config,
'parts': {k: str(v) for k, v in paths.items()},
'mix': '0.5*G probability + 0.5*Q probability; normalize roundoff only',
'kit_revision': '25c6ec7b6b5f44a216c63c3a700de3a9c6d6699d',
'policy': 'native prompts and temperatures; no item-specific rules or truncation'}
def _read(self, timeout):
"""One reply line within `timeout` seconds in total (also across a partial line). A bad line closes the backend: a stale reply must never be matched to a later request."""
deadline = time.monotonic() + timeout
fd = self.process.stdout.fileno()
while b'\n' not in self._buf:
remaining = deadline - time.monotonic()
if remaining <= 0 or not select.select([fd], [], [], remaining)[0]:
self.close()
raise TimeoutError('REAX backend response deadline exceeded; backend closed')
chunk = os.read(fd, 1 << 20)
if not chunk:
status = self.process.poll()
self.close()
raise RuntimeError('REAX backend exited: ' + str(status))
self._buf += chunk
line, _, self._buf = self._buf.partition(b'\n')
try:
reply = json.loads(line.decode('utf-8'))
except ValueError:
self.close()
raise RuntimeError('REAX backend sent an unreadable reply line; backend closed')
if not isinstance(reply, dict):
self.close()
raise RuntimeError('REAX backend sent a non-object reply line; backend closed')
return reply
# Exact reply shapes of backend.py (the 'ok' and 'error' prints at the end of its request loop).
OK_KEYS = ('Duo', 'timings_ms', 'peak_bytes', 'peak_reserved_bytes', 'optimization')
ERROR_KEYS = ('error', 'exception')
@classmethod
def _is_reply(cls, reply):
"""True only for a line shaped like one of backend.py's two replies. Known limit: a stray line with exactly the shape of a real
reply cannot be detected without request ids; only a line that is not reply-shaped closes the backend."""
status = reply.get('status')
if status == 'ok':
return all(key in reply for key in cls.OK_KEYS)
if status == 'error':
return all(key in reply for key in cls.ERROR_KEYS)
return False
def alive(self):
process = getattr(self, 'process', None)
return process is not None and process.poll() is None
def __call__(self, state, questions):
if not questions or len(questions) > 64:
raise Unsupported('supports 1..64 questions per request')
for question in questions.values():
if question['type'] not in ('choice', 'noul', 'score'):
raise Unsupported('unsupported question type ' + question['type'])
if question['type'] == 'choice' and not 2 <= len(question['criteria']) <= 255:
raise Unsupported('options per choice must be 2..255')
size = proxy_tokens(state, questions)
if size > self.max_input_proxy_tokens:
# Declared input cap: refused before any model work, never truncated.
raise Unsupported('request of %d proxy tokens exceeds this engine\'s declared maximum context length of %d proxy tokens; input is never truncated'
% (size, self.max_input_proxy_tokens))
with self.lock:
if not self.alive():
raise EngineUnavailable('REAX backend is not running (closed after a timeout, GPU failure or crash); restart the engine')
request = {'state': state, 'questions': questions,
'_kind': 'concurrent' if self.config['optimization'] != 'plain' else 'sequential',
'_ubatch': int(self.config['u_batch_size']), '_opt': self.config['optimization']}
self.process.stdin.write(json.dumps(request) + '\n')
self.process.stdin.flush()
reply = self._read(float(self.config['request_timeout']))
if not self._is_reply(reply):
# Not a reply (e.g. a stray JSON object printed by a native library): the real reply is still queued.
self.close()
raise RuntimeError('REAX backend sent a line that is not a reply; backend closed')
if reply.get('status') != 'ok':
message = reply.get('error', 'native backend error')
if reply.get('exception') in ('OutOfMemoryError', 'AcceleratorError') or 'device-side assert' in message:
self.close()
raise RuntimeError('device-side assert or GPU memory failure; engine closed: ' + message)
if reply.get('exception') == 'Unsupported' or 'maximum context length' in message:
raise Unsupported(message)
raise RuntimeError(message)
return reply['Duo'], {k: reply[k] for k in ('timings_ms', 'peak_bytes', 'peak_reserved_bytes', 'optimization')}
def warmup(self):
super().warmup()
path = self.config.get('warmup_rows')
if path:
rows = [json.loads(line) for line in Path(path).read_text().splitlines()]
fits = [row for row in rows if proxy_tokens(row['state'], row['questions']) <= self.max_input_proxy_tokens]
warm = sorted(fits, key=lambda row: len(json.dumps(row)), reverse=True)[:3] + fits[:10]
for row in warm:
try:
self(row['state'], row['questions'])
except Unsupported:
pass
def runtime(self):
return self.ready
def close(self):
process = getattr(self, 'process', None)
if process is not None:
self.process = None
self._buf = b''
if process.poll() is None:
process.stdin.close()
try:
process.wait(timeout=20)
except subprocess.TimeoutExpired:
process.terminate()
try:
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
|