File size: 7,596 Bytes
1d2de8a | 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 | """Backend-neutral contract-v3 service; restart replays immutable inputs, not tensors.
Engine protocol: identity (mapping or method) -> JSON mapping with fingerprint; prefill(parts) -> state;
ask(state, block, n, execution='cached'|'full') -> letter_logits + branch_tokens.
Engine.ask must fork its prefix cache for every branch. One service serializes all
engine calls. A service instance exclusively owns its engine.
"""
import copy
import hashlib
import json
import re
import shutil
import uuid
from pathlib import Path
import numpy as np
from solomon.service_states import Service as ReadoutService, handler
def _json(value):
return json.dumps(value, sort_keys=True, separators=(',', ':'), allow_nan=False)
def _identity(engine):
value = engine.identity() if callable(engine.identity) else engine.identity
if not isinstance(value, dict) or not isinstance(value.get('fingerprint'), str) or not value['fingerprint']:
raise ValueError('engine identity requires a nonempty fingerprint')
return json.loads(_json(value))
def _parts(document):
if isinstance(document, str):
document = [{'text': document}]
if not isinstance(document, list) or not document:
raise ValueError('document must be text or a nonempty parts list')
result = []
for part in document:
if not isinstance(part, dict) or set(part) not in ({'text'}, {'image'}):
raise ValueError('each part must contain exactly text or image')
name, value = next(iter(part.items()))
if not isinstance(value, str) or not value.strip():
raise ValueError('document parts must contain nonempty strings')
result.append({name: value})
return result
def _digest(parts):
# JSON framing preserves boundaries; two text parts cannot alias one text part.
payload = [p if 'text' in p else {'image_sha256': hashlib.sha256(Path(p['image']).read_bytes()).hexdigest()} for p in parts]
return hashlib.sha256(_json(payload).encode()).hexdigest()
class _CheckedEngine:
def __init__(self, engine):
self.backend = engine
def identity(self):
value = self.backend.identity
return value() if callable(value) else value
def prefill(self, document):
return self.backend.prefill(copy.deepcopy(document))
def ask(self, state, block, n, execution='cached'):
result = self.backend.ask(state, block, n, execution=execution)
logits = np.asarray(result.get('letter_logits'), dtype=float)
tokens = result.get('branch_tokens')
if logits.shape != (n,) or not np.isfinite(logits).all():
raise ValueError('engine returned invalid letter logits')
if type(tokens) is not int or tokens < 0:
raise ValueError('engine returned invalid branch token count')
return result
class Service(ReadoutService):
"""Shared existing readout logic with CUDA-safe persistence and identity checks."""
def __init__(self, store, engine, design=None):
design = copy.deepcopy(design or {'single_choice': 'R', 'ordered': 'R'})
allowed = {'R', 'S', 'P', 'R+avg2', 'R+avg3', 'R+avgall', 'R+debias'}
if any(design.get(k) not in allowed for k in ('single_choice', 'ordered')):
raise ValueError('unsupported readout design')
if design['ordered'] not in {'R', 'S'}:
raise ValueError('ordered design must preserve caller level order (R or S)')
if 'R+debias' in design.values():
prior = design.get('position_prior', {})
for n in range(2, 9):
values = np.asarray(prior.get(str(n+2)), dtype=float)
if values.shape != (n+2,) or not np.isfinite(values).all():
raise ValueError('debias design needs finite priors for 2 to 8 options')
super().__init__(store, _CheckedEngine(engine), design)
self.runtime_identity = _identity(self.engine)
self.design_sha256 = hashlib.sha256(_json(self.design).encode()).hexdigest()
def _check_runtime(self):
if _identity(self.engine) != self.runtime_identity:
self.state = self.state_id = None
raise ValueError('engine identity changed; create a new service instance')
def create(self, document):
parts = _parts(document)
with self.lock:
self._check_runtime()
key = uuid.uuid4().hex
pending = self.store / ('.pending-' + key)
final = self.store / key
pending.mkdir()
try:
saved = []
for i, part in enumerate(parts):
if 'text' in part:
saved.append(dict(part))
else:
source = Path(part['image'])
filename = f'image-{i}{source.suffix}'
(pending / filename).write_bytes(source.read_bytes())
saved.append({'image': str((final / filename).resolve())})
digest_parts = [p if 'text' in p else {'image': str(pending / Path(p['image']).name)} for p in saved]
record = {'schema': 'cuda-service-inputs-v1', 'state_id': key,
'document': saved, 'document_sha256': _digest(digest_parts),
'runtime_identity': self.runtime_identity,
'persistence': 'immutable_inputs_restart_reprefill'}
(pending / 'record.json').write_text(_json(record))
pending.rename(final)
except BaseException:
shutil.rmtree(pending, ignore_errors=True)
raise
return {'state_id': key, 'contract': 'solomon-answer-contract-v3',
'document_sha256': record['document_sha256'],
'persistence': record['persistence']}
def _warm(self, key):
if not isinstance(key, str) or re.fullmatch('[0-9a-f]{32}', key) is None:
raise ValueError('invalid state identifier')
self._check_runtime()
root = self.store / key
record = json.loads((root / 'record.json').read_text())
if record.get('schema') != 'cuda-service-inputs-v1' or record.get('state_id') != key:
raise ValueError('invalid saved state record')
if record.get('runtime_identity') != self.runtime_identity:
raise ValueError('saved state uses a different runtime identity')
parts = _parts(record['document'])
for part in parts:
if 'image' in part and Path(part['image']).resolve().parent != root.resolve():
raise ValueError('saved image must stay inside its state directory')
if _digest(parts) != record.get('document_sha256'):
raise ValueError('saved document content changed')
if self.state_id != key:
# Clear both before prefill: a failed prefill cannot reuse a stale state.
self.state = self.state_id = None
self.state = self.engine.prefill(parts)
self.state_id = key
return self.state
def ask(self, key, task, **kwargs):
if task not in ('boolean', 'single', 'ordered', 'multilabel', 'entity'):
raise ValueError('unknown answer type')
with self.lock:
result = super().ask(key, task, **kwargs)
result.update(runtime_identity=copy.deepcopy(self.runtime_identity),
design_sha256=self.design_sha256,
persistence='immutable_inputs_restart_reprefill')
return result
|