File size: 11,712 Bytes
068cb14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
"""Standard-library HTTP adapter. No model is loaded until main()."""
from __future__ import annotations

import argparse
import json
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

if __package__:
    from .core import ItemFailure, ServiceFailure, UnsupportedRequest, Service, ROOT, load_json, validate_config, apply_profile, validate_calibration, validate_threshold
else:
    from core import ItemFailure, ServiceFailure, UnsupportedRequest, Service, ROOT, load_json, validate_config, apply_profile, validate_calibration, validate_threshold

BINDS = ("127.0.0.1", "0.0.0.0")


def make_handler(service):
    def service_dead():
        check = getattr(service, 'is_dead', None)
        return check() if callable(check) else getattr(service, 'dead', False)

    class Handler(BaseHTTPRequestHandler):
        protocol_version='HTTP/1.1'

        def handle_expect_100(self):
            if self.command == 'POST' and service_dead():
                self.close_connection=True
                self.respond(503, {'error':'service failure'})
                return False
            try:
                n=int(self.headers.get('Content-Length','0'))
                if n<=0 or n>service.config['server']['max_request_bytes'] or self.headers.get('Transfer-Encoding'):
                    raise ValueError
            except (ValueError,TypeError):
                self.close_connection=True
                self.respond(400,{'error':'Invalid request size'})
                return False
            self.send_response_only(100)
            self.end_headers()
            return True

        def setup(self):
            super().setup()
            self.connection.settimeout(service.config["server"]["socket_timeout_seconds"])

        def log_message(self, format, *args):
            # Avoid logging request contents, model paths or network endpoints.
            pass

        def respond(self, status, value):
            data = json.dumps(value, ensure_ascii=False, allow_nan=False).encode("utf-8")
            self.send_response(status)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self):
            if self.path.split('?', 1)[0].rstrip('/') == '/health':
                if __package__:
                    from .attention import attention_health
                else:
                    from attention import attention_health
                health=service.health()
                # Attention backend and runtime versions are information, never warnings.
                info=health.setdefault('info',{})
                info.update(attention_health())
                info['runtime']=getattr(service.backend,'runtime_info',None)
                self.respond(200, health)
            elif self.path.split('?',1)[0].rstrip('/') == '/engine_stats':
                # Read-only engine counters for maintainer usage reconciliation; not used by evaluation.
                try: self.respond(200,service.engine_stats())
                except Exception: self.respond(503,{'error':'Engine counters unavailable'})
            else:
                self.respond(404, {'error':'Unknown route'})

        def do_POST(self):
            # After an unrecoverable engine error or the third consecutive item failure,
            # every POST is answered 503 without reading or processing the request.
            if service_dead():
                self.close_connection=True
                self.respond(503, {'error':'service failure'})
                return
            if self.path.split("?", 1)[0].rstrip("/") != "/v1/systemone":
                self.close_connection=True
                self.respond(404, {"error": "Unknown route"})
                return
            try:
                if self.headers.get("Transfer-Encoding"):
                    raise ValueError("Transfer-Encoding is unsupported")
                length = int(self.headers.get("Content-Length", "0"))
                if length <= 0 or length > service.config["server"]["max_request_bytes"]:
                    raise ValueError("Invalid request size")
                def reject(value):
                    raise ValueError("Non-finite JSON value")
                raw = self.rfile.read(length)
                if len(raw) != length:
                    raise ValueError("Incomplete request body")
                payload = json.loads(raw.decode("utf-8"), parse_constant=reject)
                # Validate before inference so model ValueError is never a client error.
                if __package__:
                    from .core import normalize_request
                else:
                    from core import normalize_request
                normalize_request(payload, service.config["server"]["max_questions"])
            except UnsupportedRequest:
                self.close_connection=True
                self.respond(422,{"error":"Unsupported systemone request"})
                return
            except (ValueError, TypeError, UnicodeError, OSError):
                self.close_connection=True
                self.respond(400, {"error": "Invalid systemone request"})
                return
            # Failures never produce a substitute answer: 422 for one item, 503 for the service.
            try:
                result = service.answer(payload)
                json.dumps(result, allow_nan=False)
            except ItemFailure as exc:
                print("item failure: "+type(exc.__cause__ or exc).__name__, flush=True)
                self.respond(422, {"error": "item could not be processed"})
                return
            except ServiceFailure as exc:
                print("service failure: "+str(exc), flush=True)
                self.respond(503, {"error": "service failure"})
                return
            except Exception as exc:
                print("service failure: "+type(exc).__name__, flush=True)
                self.respond(503, {"error": "service failure"})
                return
            self.respond(200, result)

    return Handler


def install_extensions(service, config):
    """Serving extensions named in config.json, installed on the running service in this order:
    item_deadline (per-decision deadline of the drafting pass, with the item rules of the
    `item_rules` block: the question types that may take the drafting pass and the low-margin choice
    re-read), then post_calibration (per-type calibration of the reported probabilities). A
    configuration without these blocks serves the unchanged service."""
    if __package__:
        from . import backends, core, item_deadline, post_calibration
    else:
        import backends, core, item_deadline, post_calibration
    if config.get('item_deadline') is not None:
        rules = config.get('item_rules') or {}
        types = rules.get('slow_types')
        service.warmup()   # the drafting speed model is timed on a warmed engine, as in the measured setup
        item_deadline.install(core, backends, service.backend, config['item_deadline'],
                              system_prompt=config['prompt']['system']).wrap_service(
            service, slow_types=None if types is None or set(types) == {'choice', 'noul', 'score'} else types,
            reread_tau=rules.get('reread_tau') or None)
    if config.get('post_calibration') is not None:
        post_calibration.wrap_service(service, config['post_calibration'], core)
    return service


class QuietHTTPServer(ThreadingHTTPServer):
    # Request workers must not keep a stopped service alive while draining.
    daemon_threads = True

    def handle_error(self, request, client_address):
        # The base implementation prints a peer endpoint on disconnected clients.
        print("HTTP connection failed", flush=True)


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "config.json")
    parser.add_argument("--model")
    parser.add_argument("--quant", choices=("int8", "bf16", "awq"))
    parser.add_argument("--adapter")
    parser.add_argument("--backend", choices=("transformers", "vllm"))
    parser.add_argument("--bind", choices=BINDS, help="Listening address; defaults to server.bind in config")
    parser.add_argument("--port", type=int)
    parser.add_argument("--mode", choices=("routed", "uniform", "fast_only"))
    parser.add_argument("--calibration", type=Path)
    parser.add_argument('--state-dir', type=Path, help='Absolute directory for run state (required)')
    parser.add_argument('--run-id')
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--model-key')
    parser.add_argument('--quota-tier',choices=('normal','hard','long'))
    parser.add_argument('--cost-fuse',choices=('on','off'))
    parser.add_argument('--no-prefix-caching',action='store_true')
    parser.add_argument('--allow-placeholder-threshold',action='store_true')
    args = parser.parse_args(argv)
    if args.state_dir is None or not args.state_dir.is_absolute():
        parser.error("--state-dir must be given as an absolute path")
    c = load_json(args.config)
    for arg, field in (("model", "path"), ("quant", "quant"), ("adapter", "adapter"), ("backend", "backend")):
        if getattr(args, arg) is not None:
            c["model"][field] = getattr(args, arg)
    for key in ("bind", "port"):
        if getattr(args, key) is not None:
            c["server"][key] = getattr(args, key)
    if args.mode is not None:
        c["routing"]["mode"] = args.mode
    if args.state_dir is not None:
        c['routing']['state_dir'] = str(args.state_dir)
    if args.run_id:
        c['routing']['run_id'] = args.run_id
    c['routing']['resume'] = args.resume
    if args.model_key:
        c['model']['model_key'] = args.model_key
    apply_profile(c)
    if __package__: from .core import apply_budget_tier
    else: from core import apply_budget_tier
    apply_budget_tier(c,args.quota_tier)
    if args.mode is not None:c["routing"]["mode"]=args.mode
    if args.quant is not None and args.quant!=c['model']['quant']:parser.error('Precision conflicts with model profile')
    if args.backend is not None and args.backend!=c['model']['backend']:parser.error('Backend conflicts with model profile')
    if args.no_prefix_caching: c["model"]["enable_prefix_caching"]=False
    if args.cost_fuse:c['routing']['cost_fuse']['enabled']=args.cost_fuse=='on'
    validate_config(c)
    validate_threshold(c,args.allow_placeholder_threshold)
    if not c["model"]["path"] or not c["server"]["bind"]:
        parser.error("Set --model and --bind, or their config fields")
    calibration = load_json(args.calibration or ROOT / c["calibration_file"])
    validate_calibration(calibration, c["model"]["model_key"])
    if __package__:
        from .backends import create_backend
        from .quota import RuntimeCounter
    else:
        from backends import create_backend
        from quota import RuntimeCounter
    counter = RuntimeCounter(c)
    service = Service(c, calibration, create_backend(c), counter, allow_placeholder_threshold=args.allow_placeholder_threshold)
    install_extensions(service, c)
    if c["server"].get("warmup", True):
        service.warmup()
    httpd = QuietHTTPServer((c["server"]["bind"], c["server"]["port"]), make_handler(service))
    print("systemone service ready", flush=True)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        httpd.server_close()


if __name__ == "__main__":
    main()