gpu-worker: fork-per-accept + VRAM-aware admission

Single-PID parent persistent listener; each accept forks a short-lived
child handler that owns one bend-cuda subprocess + responds + exits.
Linux COW handles memory; OS scheduler distributes across cores.
N concurrent requests = N children + parent — naturally VRAM-isolated.

VRAM admission: wait-vram-clear queries nvidia-smi before each fork,
blocks accept when used > *vram-budget-mib* (default 22000, override
via LUMBDA_VRAM_BUDGET_MIB env). 24G card with avg 1-2GB per bin
supports 4-12 concurrent comfortably.

Requires lumbda c-tier fork-self / waitpid-nonblock / exit-immediate /
sleep primitives (commit 81ac49e). Child uses exit-immediate not exit
to avoid dual-cleanup hang on shared parent state.
This commit is contained in:
russell@unturf.com 2026-06-11 09:25:07 -04:00
parent 81ac49ece0
commit 092a8ed741
No known key found for this signature in database

View file

@ -673,11 +673,133 @@
(tcp-close client)
#t))))))))
;;; -- VRAM-aware admission control ------------------------------
;;;
;;; bend-cuda holds ~bin-size GB of VRAM during its run (per
;;; foxhop empirical: 38 MiB baseline + bin-load roughly bin-size).
;;; To avoid OOM under fork-per-accept, query nvidia-smi before each
;;; fork and BLOCK accept when used > *vram-budget-mib*.
;;;
;;; Budget is per-host-policy; 22 GB of 24 GB leaves headroom for
;;; one in-flight bin tail + framework overhead. Override with env
;;; LUMBDA_VRAM_BUDGET_MIB.
(define *vram-budget-mib*
(let ((env (get-environment-variable "LUMBDA_VRAM_BUDGET_MIB")))
(cond
((and env (> (string-length env) 0)) (string->number env))
(else 22000))))
(define (vram-used-mib)
;; nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits
;; returns one integer line per GPU; sum if multi-GPU host.
(let* ((pair (spawn-process-stdio
"/usr/bin/nvidia-smi"
'("--query-gpu=memory.used"
"--format=csv,noheader,nounits"))))
(cond
((eq? pair #f) 0) ;; fail-open: assume 0 if nvidia-smi missing
(else
(close-port (car pair)) ;; close stdin to child
(let loop ((sum 0))
(let ((line (read-line (cdr pair))))
(cond
((eof-object? line) (close-port (cdr pair)) sum)
(else
(let ((n (string->number (string-trim line))))
(loop (+ sum (cond (n n) (else 0)))))))))))))
(define (string-trim s)
;; trim leading + trailing whitespace.
(let* ((len (string-length s))
(start (let loop ((i 0))
(cond ((>= i len) len)
((char-whitespace? (string-ref s i)) (loop (+ i 1)))
(else i))))
(end (let loop ((i (- len 1)))
(cond ((< i start) start)
((char-whitespace? (string-ref s i)) (loop (- i 1)))
(else (+ i 1))))))
(substring s start end)))
(define (wait-vram-clear)
;; Block until VRAM drops below budget. Reap zombies while waiting
;; so parent doesn't accumulate them.
(let loop ()
(waitpid-nonblock)
(let ((used (vram-used-mib)))
(cond
((< used *vram-budget-mib*) #t)
(else
(sleep 1)
(loop))))))
;;; -- entry -----------------------------------------------------
;;; run-loop — fork-per-accept pattern.
;;;
;;; Accept connection → check VRAM budget → fork child handler →
;;; parent reaps zombies + returns to accept. Child runs handle-one
;;; (which spawns bend-cuda + waits + responds) then exits — its
;;; bend-cuda subprocess holds VRAM for the duration of that one
;;; request, naturally isolated from sibling children.
;;;
;;; This gives concurrent dispatch on a single PID without lumbda
;;; needing threading primitives. Linux COW handles parent->child
;;; memory; OS scheduler distributes across cores.
(define (run-loop server)
(handle-one server)
(run-loop server))
(waitpid-nonblock) ;; reap any completed child
(wait-vram-clear) ;; block accept if GPU near full
(let ((client (tcp-accept server)))
(cond
((eq? client #f) (run-loop server))
(else
(let ((pid (fork-self)))
(cond
((eq? pid #f)
;; fork failed — fall back to serial handle
(handle-one-client client server)
(run-loop server))
((eq? pid 0)
;; child: handle this one client, then exit
(handle-one-client client server)
(exit 0))
(else
;; parent: close our copy of client fd, loop to accept
(tcp-close client)
(run-loop server))))))))
;;; handle-one-client — same dispatch as the old handle-one body but
;;; takes the already-accepted client port as arg (no accept call).
(define (handle-one-client client server)
(let ((payload (wire-recv-raw client)))
(cond
((eq? payload #f) (tcp-close client))
((and (>= (string-length payload) 4)
(string=? (substring payload 0 4) "BSHK"))
(handle-binary-shake client payload)
(tcp-close client))
((and (>= (string-length payload) 4)
(string=? (substring payload 0 4) "BCGB"))
(handle-binary-cgbn client payload)
(tcp-close client))
((and (>= (string-length payload) 4)
(string=? (substring payload 0 4) "BSCP"))
(handle-binary-secp client payload)
(tcp-close client))
((and (>= (string-length payload) 4)
(string=? (substring payload 0 4) "BSRT"))
(handle-binary-sort client payload)
(tcp-close client))
((and (>= (string-length payload) 4)
(string=? (substring payload 0 4) "BSB3"))
(handle-binary-blake3 client payload)
(tcp-close client))
(else
(let* ((req (read-from-string payload))
(resp (handle-request req)))
(wire-send client resp)
(tcp-close client))))))
;; Optional registration — only spawn the daemon when its binary is
;; reachable. Lets a worker host serve a subset of forms without