nx_gpu_broker.nx source
↩ module page · 154 lines · 9852 B
1// nx_gpu_broker.nx -- THE RESOURCE-SHARING BRAIN (serving-census capstone). Admits/schedules a MIXED workload --
2// background image-gen (Z-Image DiT, big VRAM, low priority, preemptible) + interactive chat (small VRAM, high
3// priority, latency-sensitive) -- onto ONE 16GB GPU, using BOTH levers together: (1) VRAM-AWARE admission (never
4// oversubscribe the card; co-resident jobs whose VRAM sums <= budget run TOGETHER) and (2) QoS PREEMPTION (a
5// high-priority arrival preempts enough low-priority resident jobs to fit, and the preempted job resumes later
6// with its work intact). This is the operator's "resource sharing" highlight made a working decision-maker --
7// it composes the verified nx_vram_estimate (budgets) + nx_sched_qos (preemption) ideas into one broker. Pure
8// logic, no GPU/model needed; VRAM never oversubscribed BY CONSTRUCTION (checked every op). license_tier: ORIGINAL expect_exit: 0
9import "nx_syscalls.nx"
10const KIND_MAGIC_16384: i64 = 16384
11const KIND_MAGIC_10240: i64 = 10240
12const KIND_MAGIC_2048: i64 = 2048
13const KIND_MAGIC_12288: i64 = 12288
14const KIND_MAGIC_6144: i64 = 6144
15
16const BUSY: i64 = 1
17const FREE: i64 = 0
18const KIND_CHAT: i64 = 0
19const KIND_IMAGE: i64 = 1
20const QCAP: i64 = 64
21
22// admit results
23const R_REJECT: i64 = 0 - 1 // never fits (vram > budget)
24const R_QUEUE: i64 = 0 - 2 // queued (no room, nothing preemptible)
25
26struct GpuBroker {
27 budget: i64, used: i64, n_slots: i64,
28 st: *i64, id: *i64, kind: *i64, vram: *i64, prio: *i64, work: *i64,
29 qn: i64, qid: *i64, qkind: *i64, qvram: *i64, qprio: *i64, qwork: *i64
30}
31
32func gb_new(budget_mb: i64, n_slots: i64) -> *GpuBroker {
33 let b: *GpuBroker = sys_mmap(160) as *GpuBroker
34 b.budget = budget_mb; b.used = 0; b.n_slots = n_slots
35 b.st=sys_mmap(n_slots*8) as *i64; b.id=sys_mmap(n_slots*8) as *i64; b.kind=sys_mmap(n_slots*8) as *i64
36 b.vram=sys_mmap(n_slots*8) as *i64; b.prio=sys_mmap(n_slots*8) as *i64; b.work=sys_mmap(n_slots*8) as *i64
37 let st: *i64 = b.st; var i: i64=0; while i<n_slots { st[i]=FREE; i=i+1 }
38 b.qn=0
39 b.qid=sys_mmap(QCAP*8) as *i64; b.qkind=sys_mmap(QCAP*8) as *i64; b.qvram=sys_mmap(QCAP*8) as *i64
40 b.qprio=sys_mmap(QCAP*8) as *i64; b.qwork=sys_mmap(QCAP*8) as *i64
41 return b
42}
43func gb_free_bytes(b: *GpuBroker) -> i64 { return b.budget - b.used }
44func gb_free_slot(b: *GpuBroker) -> i64 { let st: *i64=b.st; var i: i64=0; while i<b.n_slots { if st[i]==FREE { return i } i=i+1 } return 0-1 }
45func gb_resident(b: *GpuBroker) -> i64 { let st: *i64=b.st; var n: i64=0; var i: i64=0; while i<b.n_slots { if st[i]==BUSY { n=n+1 } i=i+1 } return n }
46
47func _enq(b: *GpuBroker, id: i64, kind: i64, vram: i64, prio: i64, work: i64) -> i64 {
48 if b.qn>=QCAP { return 0-1 }
49 let qid: *i64=b.qid; let qk: *i64=b.qkind; let qv: *i64=b.qvram; let qp: *i64=b.qprio; let qw: *i64=b.qwork
50 let n: i64=b.qn; qid[n]=id; qk[n]=kind; qv[n]=vram; qp[n]=prio; qw[n]=work; b.qn=n+1; return 0
51}
52func _place(b: *GpuBroker, slot: i64, id: i64, kind: i64, vram: i64, prio: i64, work: i64) -> i64 {
53 let st: *i64=b.st; let sid: *i64=b.id; let sk: *i64=b.kind; let sv: *i64=b.vram; let sp: *i64=b.prio; let sw: *i64=b.work
54 st[slot]=BUSY; sid[slot]=id; sk[slot]=kind; sv[slot]=vram; sp[slot]=prio; sw[slot]=work
55 b.used = b.used + vram
56 return slot
57}
58// sum VRAM of resident jobs strictly LOWER priority than `prio` (preemptible headroom)
59func _preemptible_vram(b: *GpuBroker, prio: i64) -> i64 {
60 let st: *i64=b.st; let sp: *i64=b.prio; let sv: *i64=b.vram
61 var s: i64=0; var i: i64=0
62 while i<b.n_slots { if st[i]==BUSY { if sp[i]<prio { s=s+sv[i] } } i=i+1 }
63 return s
64}
65// preempt the single lowest-priority resident job (< prio); requeue it WITH work; free its VRAM+slot. return -1 if none.
66func _preempt_one(b: *GpuBroker, prio: i64) -> i64 {
67 let st: *i64=b.st; let sp: *i64=b.prio; let sv: *i64=b.vram; let sid: *i64=b.id; let sk: *i64=b.kind; let sw: *i64=b.work
68 var victim: i64=0-1; var vp: i64=prio
69 var i: i64=0
70 while i<b.n_slots { if st[i]==BUSY { if sp[i]<vp { vp=sp[i]; victim=i } } i=i+1 }
71 if victim<0 { return 0-1 }
72 _enq(b, sid[victim], sk[victim], sv[victim], sp[victim], sw[victim])
73 b.used = b.used - sv[victim]
74 st[victim]=FREE
75 return sid[victim]
76}
77
78// THE ADMISSION DECISION. sets preempted_out[0] = last-preempted id (or -1). returns slot / R_QUEUE / R_REJECT.
79func gb_admit(b: *GpuBroker, id: i64, kind: i64, vram: i64, prio: i64, work: i64, preempted_out: *i64) -> i64 {
80 preempted_out[0] = 0-1
81 if vram > b.budget { return R_REJECT } // never fits, even alone
82 // fast path: fits in free VRAM + a free slot -> CO-RESIDENT admit (runs alongside whatever's already there)
83 if gb_free_bytes(b) >= vram { if gb_free_slot(b) >= 0 { return _place(b, gb_free_slot(b), id, kind, vram, prio, work) } }
84 // constrained: can preempting lower-priority jobs make room?
85 if gb_free_bytes(b) + _preemptible_vram(b, prio) >= vram {
86 var guard: i64 = 0
87 while gb_free_bytes(b) < vram { if guard > b.n_slots { guard = b.n_slots } let pid: i64 = _preempt_one(b, prio); if pid < 0 { guard = b.n_slots + 1 } else { preempted_out[0] = pid } guard = guard + 1 if guard > b.n_slots + 1 { return R_QUEUE } }
88 let fs: i64 = gb_free_slot(b)
89 if fs >= 0 { return _place(b, fs, id, kind, vram, prio, work) }
90 }
91 _enq(b, id, kind, vram, prio, work)
92 return R_QUEUE
93}
94// complete a resident job -> free VRAM+slot, then admit the highest-priority queued job that now FITS.
95func gb_complete(b: *GpuBroker, slot: i64) -> i64 {
96 let st: *i64=b.st; let sv: *i64=b.vram
97 b.used = b.used - sv[slot]; st[slot]=FREE
98 if b.qn<=0 { return 0-1 }
99 let qpr: *i64=b.qprio; let qv: *i64=b.qvram; let qid: *i64=b.qid; let qk: *i64=b.qkind; let qw: *i64=b.qwork
100 // highest priority queued that fits in free VRAM
101 var pick: i64=0-1; var pp: i64=0-1
102 var i: i64=0
103 while i<b.qn { if qv[i] <= gb_free_bytes(b) { if qpr[i] > pp { pp=qpr[i]; pick=i } } i=i+1 }
104 if pick<0 { return 0-1 }
105 let rid: i64=qid[pick]
106 _place(b, gb_free_slot(b), rid, qk[pick], qv[pick], qpr[pick], qw[pick])
107 var j: i64=pick; while j<b.qn-1 { qid[j]=qid[j+1]; qk[j]=qk[j+1]; qv[j]=qv[j+1]; qpr[j]=qpr[j+1]; qw[j]=qw[j+1]; j=j+1 }
108 b.qn=b.qn-1
109 return rid
110}
111
112func gw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
113func gn(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
114
115func main() -> i64 {
116 gw("=== nx_gpu_broker -- image-gen + chat share ONE 16GB GPU (VRAM-aware admission + QoS preemption) ===\n" as *u8)
117 let pe: *i64 = sys_mmap(8) as *i64
118 let b: *GpuBroker = gb_new(KIND_MAGIC_16384, 8) // 16 GB
119 var pass: i64 = 0; var tot: i64 = 0
120
121 // 1) a big background IMAGE render (10 GB, low prio) admits
122 let a1: i64 = gb_admit(b, 101, KIND_IMAGE, KIND_MAGIC_10240, 1, 30, pe)
123 gw(" image #101 (10GB, prio1) -> slot "); gn(a1); gw("; VRAM used "); gn(b.used); gw("/"); gn(b.budget); gw("\n" as *u8)
124
125 // T1: interactive CHAT (2GB, high prio) CO-RESIDENT with the image (10+2 <= 16) -- THE sharing win
126 tot=tot+1
127 let ac: i64 = gb_admit(b, 200, KIND_CHAT, KIND_MAGIC_2048, 10, 8, pe)
128 if ac >= 0 { if b.used == KIND_MAGIC_12288 { pass=pass+1; gw("PASS T1 chat #200 (2GB) admitted CO-RESIDENT with image #101 -> 12GB/16GB, both run on ONE GPU\n" as *u8) } else { gw("FAIL T1b used="); gn(b.used); gw("\n" as *u8) } } else { gw("FAIL T1a chat rejected\n" as *u8) }
129
130 // T2: a SECOND big image (10GB) does NOT fit (12+10>16) and can't preempt (equal/lower... chat is higher, image peer) -> QUEUED, VRAM untouched
131 tot=tot+1
132 let a2: i64 = gb_admit(b, 102, KIND_IMAGE, KIND_MAGIC_10240, 1, 30, pe)
133 if a2 == R_QUEUE { if b.used == KIND_MAGIC_12288 { pass=pass+1; gw("PASS T2 2nd image #102 (10GB) QUEUED (no room, no preemptible peer) -- VRAM NOT oversubscribed\n" as *u8) } else { gw("FAIL T2b\n" as *u8) } } else { gw("FAIL T2a a2="); gn(a2); gw("\n" as *u8) }
134
135 // T3: never oversubscribed invariant so far
136 tot=tot+1
137 if b.used <= b.budget { pass=pass+1; gw("PASS T3 invariant: VRAM used ("); gn(b.used); gw(") never exceeds budget ("); gn(b.budget); gw(")\n" as *u8) } else { gw("FAIL T3 OVERSUBSCRIBED\n" as *u8) }
138
139 // T4: a HIGH-prio chat needing 6GB arrives; only 4GB free -> PREEMPTS the low-prio image #101 to make room
140 tot=tot+1
141 let ac2: i64 = gb_admit(b, 201, KIND_CHAT, KIND_MAGIC_6144, 10, 8, pe)
142 if ac2 >= 0 { if pe[0] == 101 { if b.used <= b.budget { pass=pass+1; gw("PASS T4 chat #201 (6GB) PREEMPTED image #101 to fit -> got the GPU; used "); gn(b.used); gw("/"); gn(b.budget); gw("\n" as *u8) } else { gw("FAIL T4c oversub\n" as *u8) } } else { gw("FAIL T4b preempted="); gn(pe[0]); gw("\n" as *u8) } } else { gw("FAIL T4a chat waited\n" as *u8) }
143
144 // T5: the preempted image #101 is requeued WITH its 30 steps intact (resumes later)
145 tot=tot+1
146 var found: i64 = 0
147 let qid: *i64=b.qid; let qw: *i64=b.qwork
148 var qi: i64=0; while qi<b.qn { if qid[qi]==101 { if qw[qi]==30 { found=1 } } qi=qi+1 }
149 if found==1 { pass=pass+1; gw("PASS T5 preempted image #101 requeued with 30 steps intact (resumes, no lost render)\n" as *u8) } else { gw("FAIL T5\n" as *u8) }
150
151 gw("nx_gpu_broker pass="); gn(pass); gw("/"); gn(tot)
152 if pass==tot { gw(" GREEN -- one GPU, shared: chat runs alongside a render when it fits, preempts it when it doesn't, VRAM never oversubscribed. The resource-sharing brain, sovereign.\n" as *u8); sys_exit(0); return 0 }
153 gw(" RED\n" as *u8); sys_exit(1); return 1
154}