File size: 4,430 Bytes
a21d4da | 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 | """Local System One server for Drex v2.0.
python serve.py --port 8000
POST /v1/systemone accepts one decision, or ``requests``: a list of decisions.
Several decisions that fit share one forward. This process scores the pointer head
of the checkpoint.
"""
import argparse
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from inference import decide, decide_batch, load
def reject_nonfinite_json(value):
raise ValueError(f"non-finite JSON number: {value}")
def make_handler(tok, model, model_name):
lock = threading.Lock()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def log_message(self, fmt, *args):
print("%s %s" % (self.address_string(), fmt % args))
def _send(self, status, payload):
body = json.dumps(payload).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
if self.close_connection:
self.send_header("Connection", "close")
self.end_headers()
self.wfile.write(body)
def do_GET(self):
if self.path.rstrip("/") == "/health":
self._send(200, {"status": "ok", "model": model_name})
return
self._send(404, {"error": "not found"})
def do_POST(self):
if self.path.rstrip("/") != "/v1/systemone":
self._send(404, {"error": "not found"})
return
try:
length = int(self.headers.get("Content-Length", "0"))
except ValueError:
self.close_connection = True
self._send(400, {"error": "Content-Length must be an integer"})
return
if length <= 0 or length > 1_000_000:
self.close_connection = True
self._send(400, {"error": "request body must be between 1 byte and 1 MB"})
return
try:
request = json.loads(self.rfile.read(length), parse_constant=reject_nonfinite_json)
if not isinstance(request, dict):
raise ValueError("request must be an object")
if "model" in request and not isinstance(request["model"], str):
raise ValueError("model must be a string")
with lock:
if "requests" in request:
if "state" in request or "questions" in request:
raise ValueError("a batch cannot also contain state or questions")
items = request["requests"]
if not isinstance(items, list) or not items:
raise ValueError("requests must be a non-empty list")
if len(items) > 1024:
raise ValueError("requests exceeds 1024 items")
results = decide_batch(tok, model, items)
payload = {"model": request.get("model") or model_name, "results": results}
else:
payload = {"model": request.get("model") or model_name, **decide(tok, model, request)}
except ValueError as exc:
self._send(400, {"error": str(exc)})
return
except Exception as exc:
self.log_error("System One scoring failed: %r", exc)
self._send(500, {"error": "internal scoring error"})
return
self._send(200, payload)
return Handler
def main():
parser = argparse.ArgumentParser(description="Serve Drex v2.0 at POST /v1/systemone.")
parser.add_argument("--model", default=str(Path(__file__).resolve().parent), help="Checkpoint directory (default: this directory)")
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--name", default="drex-v2.0")
args = parser.parse_args()
name = args.name
tok, model = load(args.model)
server = ThreadingHTTPServer((args.host, args.port), make_handler(tok, model, name))
print(f"{name} listening on http://{args.host}:{args.port}/v1/systemone")
server.serve_forever()
if __name__ == "__main__":
main()
|