lumbda/examples/cuda-fanout/bend.lsp

;;; bend.lsp -- Lisp-smart GPU dispatch primitive.
;;;
;;; (bend expr)
;;;   inspects expr at runtime, routes to a registered GPU worker
;;;   when the cost estimator says it's worth the trip, else
;;;   evaluates locally in the original lexical scope.
;;;
;;;   The decision is per-call, based on what's in the args -- not
;;;   just the op name. Tiny SHAKE on 5 inputs stays local; big
;;;   SHAKE on 5,000,000 inputs bends to a GPU worker.
;;;
;;; (bend! expr)
;;;   forces GPU dispatch; raises an error if no worker available.
;;;
;;; Wire protocol: length-prefixed S-exp (see wire.lsp). Same shape
;;; lumbda's fleet workers already speak, so bend reuses the same
;;; tcp-* + read-from-string + write-to-string primitives every
;;; lumbda tier already provides.
;;;
;;; Per-tier integration is now ZERO new primitives -- bend.lsp
;;; works on Python, C, and asm tiers as-is.

(load "wire.lsp")

;;; -- portable error reporter -----------------------------------
;;;
;;; asm tier doesn't have `error` as a Scheme-callable builtin; the
;;; Python and C tiers do. This wraps both so the rest of the file
;;; stays portable. Displays the message, then either raises (on tiers
;;; that have it) or returns a tagged failure value.

(define (bend-error msg . rest)
  (display ";;; bend ERROR: ") (display msg)
  (for-each (lambda (x) (display " ") (display x)) rest)
  (newline)
  'bend-failure)

;;; -- tuning constants ------------------------------------------

;; Per-spawn cuda init ~200 ms; daemon-warm RPC ~100 us. Set high
;; if your workers haven't started their daemon pools yet.
(define *bend-overhead-ns* 100000)        ; 100 us once daemons are warm

;; Host throughput for SHAKE256 (cycles ? ns/byte). Comparing
;; (cost-fn args) against *bend-overhead-ns* picks the route.
(define *bend-host-ns-per-byte* 2)        ; ~500 MB/s on a modern x86

;;; -- registry -- what `bend` knows how to ship ------------------
;;;
;;; alist of (op-name . cost-fn) where cost-fn takes the literal
;;; arg list and returns estimated host-side runtime in nanoseconds.

(define *bend-gpu-ops* '())

(define (bend-register-op! name cost-fn)
  (set! *bend-gpu-ops* (cons (cons name cost-fn) *bend-gpu-ops*)))

;; Cost estimator for (cuda-shake-fanout inputs out-bytes).

;; Sum of input byte lengths x ns-per-byte.

(define (bend-shake-cost args)
  (let* ((inputs (car args))
         (n (length inputs))
         (avg-len (if (= n 0) 0
                      (quotient (string-length (car inputs)) 2))))
    (* n avg-len *bend-host-ns-per-byte*)))

(bend-register-op! 'cuda-shake-fanout bend-shake-cost)

;;; -- worker endpoint -------------------------------------------

(define *bend-worker-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)
(define *bend-worker-port* 8320)

;; Override our default localhost:8320 endpoint.

(define (bend-set-worker! host port)
  (set! *bend-worker-host* host)
  (set! *bend-worker-port* port))

;;; -- multi-worker fleet ----------------------------------------
;;;
;;; A list of (host . port) pairs. When non-empty, bend-dispatch-to-gpu
;;; rotates round-robin across our cluster; a failed pick falls forward
;;; to a peer worker on a tcp-connect error. When empty, dispatcher
;;; falls back to a single *bend-worker-host* / *bend-worker-port* pair
;;; (full back-compat with single-host callers).
;;;
;;; Set via (bend-set-workers! '(("3090-ai.foxhop.net" . 8320)))
;;; or environment variable BEND_WORKERS="host:port,host:port".
;;;
;;; Production default: 3090-ai only. ai.foxhop.net (4090) is
;;; reserved for qwen LLM; we do NOT add it to the round-robin by
;;; default. Multi-host fan-out gets enabled per long-running
;;; parallel workload — caller opts in explicitly via
;;; bend-set-workers! or BEND_WORKERS env.

(define *bend-workers* '())
(define *bend-rr-idx* 0)

(define (bend-set-workers! lst)
  (set! *bend-workers* lst)
  (set! *bend-rr-idx* 0)
  (set! *worker-health* '()))

;;; -- worker-health heartbeat -----------------------------------
;;;
;;; *worker-health* is an alist keyed by "host:port" -> a record
;;;   (last-checked-ms status vram-mb)
;;; status ∈ {ok, down}
;;;
;;; bend-pick-worker filters workers whose cached status=down sits
;;; within *health-down-cooldown-ms*, falls through to a fresh
;;; (health) probe whenever a cache entry ages past
;;; *health-cache-ttl-ms*, & ranks healthy peers by free VRAM
;;; descending so we route long-running jobs to whichever box
;;; carries the most spare GPU memory.

(define *worker-health* '())
(define *health-cache-ttl-ms* 5000)
(define *health-down-cooldown-ms* 30000)
(define *health-probe-timeout-ms* 2000)

(define (worker-key w)
  (string-append (car w) ":" (number->string (cdr w))))

(define (worker-health-get w)
  (let ((e (assoc (worker-key w) *worker-health*)))
    (cond ((eq? e #f) #f) (else (cdr e)))))

(define (worker-health-set! w status vram-mb)
  (let ((key (worker-key w))
        (entry (list (current-time-ms) status vram-mb)))
    (set! *worker-health*
          (cons (cons key entry)
                (let drop ((rest *worker-health*))
                  (cond
                    ((null? rest) '())
                    ((string=? (car (car rest)) key) (drop (cdr rest)))
                    (else (cons (car rest) (drop (cdr rest))))))))))

(define (worker-mark-down! w)
  (worker-health-set! w 'down 0))

(define (worker-mark-ok! w vram-mb)
  (worker-health-set! w 'ok vram-mb))

;; Extract (vram-free-mb N) from a (health) response form.
(define (health-find-vram form-list)
  (let loop ((rest form-list))
    (cond
      ((null? rest) 0)
      ((and (pair? (car rest)) (eq? (car (car rest)) 'vram-free-mb)
            (pair? (cdr (car rest))) (number? (car (cdr (car rest)))))
       (car (cdr (car rest))))
      (else (loop (cdr rest))))))

;; Send (health) over a fresh tcp-connect. Returns vram-free-mb on
;; success (a number, possibly 0), or symbol 'transport-fail on
;; failure. Lumbda's (eq? 0 #f) reads #t — we deliberately do NOT
;; signal failure via #f so a healthy worker reporting 0 free VRAM
;; (no GPU on host, or VRAM fully consumed by a co-tenant) still
;; ranks as up; calling code checks with `number?`.
;;
;; Backward-compat: a worker without the health handler returns
;; (error (unknown-op health)); treat that as ok with vram=0 so
;; older builds still rank as available.
(define (worker-probe-health-body w)
  (let ((sock (tcp-connect (car w) (cdr w))))
    (cond
      ((eq? sock #f) 'transport-fail)
      (else
       (wire-send sock '(health))
       (let ((reply (wire-recv sock)))
         (tcp-close sock)
         (cond
           ((eq? reply #f) 'transport-fail)
           ((not (pair? reply)) 'transport-fail)
           ((eq? (car reply) 'ok)
            (health-find-vram (cdr reply)))
           ((and (eq? (car reply) 'error)
                 (pair? (cdr reply))
                 (pair? (car (cdr reply)))
                 (eq? (car (car (cdr reply))) 'unknown-op))
            0)
           (else 'transport-fail)))))))

;; tcp-connect raises a Python ConnectionRefusedError / gaierror on
;; dead-host probes. lumbda's `guard` only catches LispErr; we use
;; with-exception-handler which catches both LispErr & Python
;; Exception. Net effect: probing a single dead worker never
;; aborts iteration over our fleet.
(define (worker-probe-health w)
  (with-exception-handler
    (lambda (exn) 'transport-fail)
    (lambda () (worker-probe-health-body w))))

(define (worker-cache-fresh? entry now)
  (let ((age (- now (car entry)))
        (status (car (cdr entry))))
    (cond
      ((eq? status 'ok) (< age *health-cache-ttl-ms*))
      ((eq? status 'down) (< age *health-down-cooldown-ms*))
      (else #f))))

;; #t if worker is reachable & not in down-cooldown. Refreshes cache
;; lazily when a probe is needed.
(define (worker-healthy? w)
  (let* ((entry (worker-health-get w))
         (now (current-time-ms)))
    (cond
      ((and entry (worker-cache-fresh? entry now))
       (eq? (car (cdr entry)) 'ok))
      (else
       (let ((vram (worker-probe-health w)))
         (cond
           ((number? vram) (worker-mark-ok! w vram) #t)
           (else (worker-mark-down! w) #f)))))))

(define (worker-cached-vram w)
  (let ((entry (worker-health-get w)))
    (cond ((eq? entry #f) 0) (else (car (cdr (cdr entry)))))))

;; In-place insertion sort by cached-vram descending.
(define (workers-rank-by-vram lst)
  (let outer ((rest lst) (acc '()))
    (cond
      ((null? rest) acc)
      (else
       (let ((w (car rest)))
         (outer (cdr rest)
                (let inner ((sorted acc))
                  (cond
                    ((null? sorted) (list w))
                    ((>= (worker-cached-vram w)
                         (worker-cached-vram (car sorted)))
                     (cons w sorted))
                    (else (cons (car sorted) (inner (cdr sorted))))))))))))

(define (workers-filter-healthy lst)
  (let loop ((rest lst) (acc '()))
    (cond
      ((null? rest) (reverse acc))
      ((worker-healthy? (car rest))
       (loop (cdr rest) (cons (car rest) acc)))
      (else (loop (cdr rest) acc)))))

;; Pick the healthiest worker (highest free VRAM). If every worker
;; is in down-cooldown, fall back to round-robin so a recovering
;; box gets a real attempt after cooldown — better than refusing
;; to dispatch at all.
(define (bend-pick-worker)
  (let* ((healthy (workers-filter-healthy *bend-workers*))
         (n (length healthy)))
    (cond
      ((= n 0)
       (let* ((all-n (length *bend-workers*))
              (idx (if (= all-n 0) 0 (remainder *bend-rr-idx* all-n))))
         (set! *bend-rr-idx* (+ *bend-rr-idx* 1))
         (list-ref *bend-workers* idx)))
      ((= n 1) (car healthy))
      (else (car (workers-rank-by-vram healthy))))))

;; Parse "host:port,host:port" into a list of (host . port) pairs. Skips
;; malformed entries silently rather than raising — keeps boot path resilient.
(define (bend-parse-workers-env s)
  (let loop ((rest s) (acc '()) (cur ""))
    (cond
      ((= (string-length rest) 0)
       (let ((parsed (bend-parse-one-worker cur)))
         (reverse (if parsed (cons parsed acc) acc))))
      ((string=? (substring rest 0 1) ",")
       (let ((parsed (bend-parse-one-worker cur)))
         (loop (substring rest 1 (string-length rest))
               (if parsed (cons parsed acc) acc) "")))
      (else
       (loop (substring rest 1 (string-length rest)) acc
             (string-append cur (substring rest 0 1)))))))

(define (bend-parse-one-worker s)
  ;; Split on first ':'; return (host . port-number) or #f if malformed.
  (let loop ((i 0))
    (cond
      ((>= i (string-length s)) #f)
      ((string=? (substring s i (+ i 1)) ":")
       (let ((host (substring s 0 i))
             (port (string->number (substring s (+ i 1) (string-length s)))))
         (cond
           ((or (= (string-length host) 0) (eq? port #f)) #f)
           (else (cons host port)))))
      (else (loop (+ i 1))))))

;; Probe TCP connect to the next pick OR the legacy single host.
;; Returns #t/#f without raising.

(define (bend-worker-available?)
  (let ((target (cond
                  ((null? *bend-workers*)
                   (cons *bend-worker-host* *bend-worker-port*))
                  (else (car *bend-workers*))))) ; cheap reachability probe
    (let ((sock (tcp-connect (car target) (cdr target))))
      (cond
        ((eq? sock #f) #f)
        (else (tcp-close sock) #t)))))

;; Open a fresh connection, send our form framed, read framed reply,

;; close. Returns our result from the worker, or raises if the worker

;; responded with (bend-error...).

(define (bend-dispatch-to-gpu quoted-form)
  (let* ((target (cond
                   ((null? *bend-workers*)
                    (cons *bend-worker-host* *bend-worker-port*))
                   (else (bend-pick-worker))))
         (sock (tcp-connect (car target) (cdr target))))
    (cond
      ((eq? sock #f)
       ;; Worker unreachable — mark down so the next pick skips it
       ;; for *health-down-cooldown-ms* before re-probing.
       (cond ((not (null? *bend-workers*)) (worker-mark-down! target)))
       (bend-error "bend: tcp-connect failed to" target))
      (else
       (wire-send sock quoted-form)
       (let ((reply (wire-recv sock)))
         (tcp-close sock)
         (cond
           ((eq? reply #f)
            (cond ((not (null? *bend-workers*)) (worker-mark-down! target)))
            (bend-error "bend: worker closed connection"))
           ((not (pair? reply)) (bend-error "bend: malformed reply" reply))
           ((eq? (car reply) 'ok) (car (cdr reply)))
           ((eq? (car reply) 'error)
            (bend-error "bend worker error:" (cdr reply)))
           (else (bend-error "bend: unexpected reply" reply))))))))

;; Optional boot-time hook: if BEND_WORKERS is set in env, parse it now.
;; Safe to call repeatedly; a missing var is a no-op.
(define (bend-load-workers-from-env!)
  (let ((s (get-environment-variable "BEND_WORKERS")))
    (cond
      ((or (eq? s #f) (= (string-length s) 0)) #f)
      (else
       (let ((lst (bend-parse-workers-env s)))
         (cond
           ((null? lst) #f)
           (else
            (bend-set-workers! lst)
            (display ";;; bend: loaded ") (display (length lst))
            (display " workers from BEND_WORKERS") (newline)
            #t)))))))

;;; -- core dispatcher -------------------------------------------

;; thunk: 0-arg lambda that evaluates the form in its original scope.

;; quoted-form: literal sexp for inspection / serialization.

;; force-gpu?: when #t, must dispatch to GPU or error.

(define (bend-dispatch thunk quoted-form force-gpu?)
  (let* ((head (if (pair? quoted-form) (car quoted-form) #f))
         (entry (assoc head *bend-gpu-ops*)))
    (cond
      (force-gpu?
       (if (and entry (bend-worker-available?))
           (bend-dispatch-to-gpu quoted-form)
           (bend-error "bend!: no GPU worker available for op" head)))
      ((not entry) (thunk))
      (else
       (let* ((cost-fn (cdr entry))
              (args (cdr quoted-form))
              (est-host-ns (cost-fn args)))
         (cond
           ((< est-host-ns *bend-overhead-ns*) (thunk))
           ((bend-worker-available?) (bend-dispatch-to-gpu quoted-form))
           (else
            (display "bend: worker unreachable, running local: ")
            (display head) (newline)
            (thunk))))))))

;;; -- function form (works on every tier -- no macros needed) ----
;;;
;;; (bend-call '(cuda-shake-fanout '("a" "b") 32))
;;;
;;; Useful from lumbda's asm tier which doesn't have define-syntax.
;;; The macro form below is a thin wrapper over this.

(define (bend-call quoted-form)
  ;; The thunk fallback returns 'no-local-fallback so callers can
  ;; pattern-match without needing the asm tier to support (bend-error...).
  (bend-dispatch (lambda () 'no-local-fallback) quoted-form #f))

(define (bend!-call quoted-form)
  (bend-dispatch (lambda () 'gpu-required-but-no-worker) quoted-form #t))

;;; -- macros live in bend-macros.lsp (Python/C tiers) ----------
;;; asm tier doesn't have define-syntax; load bend-macros.lsp only on
;;; tiers that support it. The function-form bend-call / bend!-call
;;; above is the asm-portable path.

;;; -- demo ------------------------------------------------------

(define (make-input n)
  (let loop ((i 0) (acc '()))
    (if (= i n) acc
        (loop (+ i 1) (cons "deadbeefcafebabe1234567890abcdef" acc)))))

(define (demo)
  ;; 3 inputs: cost is tiny (96 bytes), stays local.
  (display "tiny (3 inputs): bend chooses ")
  (let ((r (bend (cuda-shake-fanout '("00" "01" "deadbeef") 32))))
    (display "-> ") (display (length r)) (display " result(s)\n"))
  ;; 100k inputs x 16 bytes = 1.6 MB; should bend if worker is up.
  (display "heavy (100k inputs): bend chooses ")
  (let ((r (bend (cuda-shake-fanout (make-input 100000) 32))))
    (display "-> ") (display (length r)) (display " result(s)\n")))

;; (demo)  ; uncomment after launching gpu-worker.lsp on port 8320 (BEND)