60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
"""Stub node HTTP server for the Distributed Inference Network."""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
|
|
|
|
class _StubHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def do_POST(self):
|
|
if self.path == "/v1/infer":
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
body = json.loads(self.rfile.read(length))
|
|
messages = body.get("messages", [])
|
|
last_content = messages[-1]["content"] if messages else ""
|
|
payload = json.dumps({"text": f"stub response to: {last_content}"}).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
|
|
class StubNodeServer:
|
|
"""Single-process stub node that returns fixed inference responses."""
|
|
|
|
def __init__(self, host: str = "127.0.0.1", port: int = 0):
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._server: http.server.HTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
def start(self) -> int:
|
|
if self._server is not None:
|
|
raise RuntimeError("StubNodeServer is already running")
|
|
|
|
self._server = http.server.HTTPServer((self._host, self._requested_port), _StubHandler)
|
|
self.port = self._server.server_address[1]
|
|
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
|
self._thread.start()
|
|
return self.port
|
|
|
|
def stop(self) -> None:
|
|
if self._server is None:
|
|
return
|
|
|
|
self._server.shutdown()
|
|
self._server.server_close()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=1)
|
|
self._server = None
|
|
self._thread = None
|
|
self.port = None
|