lumbda/examples/cuda-fanout/mock-worker.py

"""mock-worker.py — Python stand-in for gpu-worker.lsp.

Listens on TCP, speaks lumbda's length-prefixed S-expression protocol
(8 ASCII digits header + payload). Dispatches `cuda-shake-fanout`
requests to the local daemon binary, returns the result.

Used until lumbda gains a `spawn-process-stdio` primitive so the pure
Scheme gpu-worker.lsp can replace this. The wire side talks to bend.lsp
unchanged."""

import os
import re
import socket
import struct
import subprocess
import sys
import threading
import time

HOST = "127.0.0.1"
# Port 8320 — BEND mnemonic:
#   8 ~= B (implied infinity B flattened; bake a cake; baby & me)
#   3 ~= E (backward)
#   2 ~= N (pivoted 90 degrees)
#   0 ~= D (flattened)
PORT = 8320
BINARY = sys.argv[1] if len(sys.argv) > 1 else "./shake256-fanout"

# spawn a warm daemon once
proc = subprocess.Popen([BINARY, "--daemon"],
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
ready = proc.stdout.readline().strip()
assert ready == "ready", f"daemon ready={ready!r}"
daemon_lock = threading.Lock()


def recv_exact(sock, n):
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


def wire_recv(sock):
    hdr = recv_exact(sock, 8)
    if hdr is None:
        return None
    plen = int(hdr.decode("ascii"))
    payload = recv_exact(sock, plen)
    return payload.decode("utf-8") if payload else None


def wire_send(sock, text):
    payload = text.encode("utf-8")
    hdr = f"{len(payload):08d}".encode("ascii")
    sock.sendall(hdr + payload)


def parse_sexp_inputs(text):
    """Two accepted shapes:
       1. Portal-style: (cuda-shake-fanout (output-bytes N) (inputs "a" "b" ...))
       2. Bend call:    (cuda-shake-fanout (quote ("a" "b" ...)) N)
       Returns (out_bytes, [hex inputs])."""
    # Shape 2 — bend call form (quoted list of strings, out-bytes integer)
    m_bend = re.match(
        r"\(cuda-shake-fanout\s+\(quote\s+\((.*?)\)\)\s+(\d+)\s*\)",
        text.strip(), re.DOTALL)
    if m_bend:
        inputs = re.findall(r'"([^"]*)"', m_bend.group(1))
        return int(m_bend.group(2)), inputs
    # Shape 1 — portal-style
    m_ob = re.search(r"\(output-bytes\s+(\d+)\)", text)
    out_bytes = int(m_ob.group(1)) if m_ob else 32
    m_in = re.search(r"\(inputs\b(.*?)\)\)", text, re.DOTALL)
    if not m_in:
        return out_bytes, []
    return out_bytes, re.findall(r'"([^"]*)"', m_in.group(1))


def dispatch_shake(inputs_hex, out_bytes):
    in_p = f"/tmp/mock-bend-in-{os.getpid()}-{int(time.time()*1e6)}.bin"
    out_p = in_p.replace("-in-", "-out-")
    # write binary input — much faster than hex portal
    with open(in_p, "wb") as fh:
        fh.write(struct.pack("<II", out_bytes, len(inputs_hex)))
        for h in inputs_hex:
            b = bytes.fromhex(h)
            fh.write(struct.pack("<I", len(b)))
            fh.write(b)
    # call daemon
    with daemon_lock:
        proc.stdin.write(f"process-bin {in_p} {out_p}\n")
        proc.stdin.flush()
        line = proc.stdout.readline().strip()
    if not line.startswith("done"):
        os.unlink(in_p)
        return None
    with open(out_p, "rb") as fh:
        n, ob = struct.unpack("<II", fh.read(8))
        hashes = [fh.read(ob).hex() for _ in range(n)]
    os.unlink(in_p); os.unlink(out_p)
    return hashes


def handle_request(text):
    # parse op head
    text = text.strip()
    if text.startswith("(cuda-shake-fanout"):
        out_bytes, inputs = parse_sexp_inputs(text)
        hashes = dispatch_shake(inputs, out_bytes)
        if hashes is None:
            return '(error "daemon failed")'
        return f"(ok ({' '.join(f'#x{h}' for h in hashes)}))"
    if text.startswith("(ping"):
        return "(ok pong)"
    return '(error "unknown-op")'


def handle_client(client):
    try:
        req = wire_recv(client)
        if req is None:
            return
        resp = handle_request(req)
        wire_send(client, resp)
    finally:
        client.close()


def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(64)
    print(f"mock-worker listening on {HOST}:{PORT}, daemon binary: {BINARY}")
    try:
        while True:
            client, _ = s.accept()
            handle_client(client)
    except KeyboardInterrupt:
        pass
    finally:
        proc.stdin.write("quit\n"); proc.stdin.flush()
        proc.stdout.readline(); proc.wait()
        s.close()


if __name__ == "__main__":
    main()