Unify continuous batching + heterogeneous runtime: decode batching, physical-core planning, disjoint VRAM/RAM placement, topp-policy warning (CPU-validated, CUDA on 6x5090) (#68)
* Fuse CUDA expert MLP execution * Group CUDA expert transfers by device * Instrument grouped CUDA expert execution * Bound grouped CUDA decode scratch * Execute expert groups across GPUs in parallel * Release host backing for multi-GPU experts * Define quality-preserving memory policies * Overlap cold expert loading with resident compute * Adapt expert placement with session LFRU * Fuse q4 expert gate and up dispatch * Plan CPU work on physical cores * Batch grouped expert CUDA kernels * Separate VRAM and RAM expert placement * Add ragged multi-sequence decode forward * feat(runtime): add continuous decode scheduler * Route concurrent API requests through batch scheduler * Harden multiplex request lifecycle and framing * Cancel disconnected multiplex requests * Bind API port before starting the engine * fix automatic KV slot allocation * add native int4 Tensor Core grouped GEMM * add Tensor Core throughput benchmark * optimize packed int4 low-row kernels * add asynchronous CUDA staging streams * document validated six-GPU dense acceleration * tune six-GPU expert hot set * raise validated expert hot-set target * add CUDA MLA absorption core * fuse grouped expert gate and up projections * Warn for explicit lossy routing flags
This commit is contained in:
ZacharyZcR
@@ -0,0 +1,45 @@
|
||||
#include "../backend_cuda.h"
|
||||
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
static double run(ColiCudaTensor *g,ColiCudaTensor *u,ColiCudaTensor *d,
|
||||
const float *x,float *y,int rows,int iterations,int mode){
|
||||
ColiCudaTensor *gs[1]={g},*us[1]={u},*ds[1]={d}; int rs[1]={rows};
|
||||
if(mode==2){setenv("COLI_CUDA_TC_INT4","1",1);setenv("COLI_CUDA_TC_MIN_ROWS","1",1);}
|
||||
else unsetenv("COLI_CUDA_TC_INT4");
|
||||
setenv("COLI_CUDA_W4_PACKED",mode==0?"0":"1",1);
|
||||
if(!coli_cuda_expert_group(gs,us,ds,rs,1,y,x))std::exit(2);
|
||||
auto begin=std::chrono::steady_clock::now();
|
||||
for(int i=0;i<iterations;i++)if(!coli_cuda_expert_group(gs,us,ds,rs,1,y,x))std::exit(2);
|
||||
auto end=std::chrono::steady_clock::now();
|
||||
return std::chrono::duration<double,std::milli>(end-begin).count()/iterations;
|
||||
}
|
||||
|
||||
int main(){
|
||||
constexpr int D=6144,I=2048,O=8;
|
||||
int device=0;if(!coli_cuda_init(&device,1))return 77;
|
||||
std::vector<unsigned char> hidden((size_t)I*D/2),down((size_t)D*I/2);
|
||||
std::vector<float> hs(I),ds(D),x((size_t)O*D),a((size_t)O*D),b((size_t)O*D),c((size_t)O*D);
|
||||
for(size_t i=0;i<hidden.size();i++)hidden[i]=(unsigned char)((i*17+29)&255);
|
||||
for(size_t i=0;i<down.size();i++)down[i]=(unsigned char)((i*13+41)&255);
|
||||
for(int i=0;i<I;i++)hs[i]=0.006f+(i%11)*0.0002f;
|
||||
for(int i=0;i<D;i++)ds[i]=0.006f+(i%7)*0.0002f;
|
||||
for(size_t i=0;i<x.size();i++)x[i]=std::sin((float)(i+1)*0.013f)*2.f;
|
||||
ColiCudaTensor *g=nullptr,*u=nullptr,*d=nullptr;
|
||||
if(!coli_cuda_tensor_upload(&g,hidden.data(),hs.data(),2,D,I,device)||
|
||||
!coli_cuda_tensor_upload(&u,hidden.data(),hs.data(),2,D,I,device)||
|
||||
!coli_cuda_tensor_upload(&d,down.data(),ds.data(),2,I,D,device))return 2;
|
||||
for(int rows: {1,2,4,8}){
|
||||
double scalar=run(g,u,d,x.data(),a.data(),rows,3,0);
|
||||
double packed=run(g,u,d,x.data(),b.data(),rows,3,1);
|
||||
double tc=run(g,u,d,x.data(),c.data(),rows,3,2);
|
||||
double pe=0,te=0,ref=0;for(int i=0;i<rows*D;i++){double p=b[i]-a[i],t=c[i]-a[i];pe+=p*p;te+=t*t;ref+=(double)a[i]*a[i];}
|
||||
std::printf("rows=%d scalar_ms=%.3f packed_ms=%.3f packed_speedup=%.3fx packed_rms=%.7f tensor_ms=%.3f tensor_speedup=%.3fx tensor_rms=%.5f\n",
|
||||
rows,scalar,packed,scalar/packed,std::sqrt(pe/(ref+1e-20)),tc,scalar/tc,std::sqrt(te/(ref+1e-20)));
|
||||
}
|
||||
coli_cuda_tensor_free(g);coli_cuda_tensor_free(u);coli_cuda_tensor_free(d);coli_cuda_shutdown();
|
||||
}
|
||||
@@ -15,6 +15,12 @@ static int close_enough(const float *got, const float *want, int n) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int relative_rms(const float *got,const float *want,int n,float limit){
|
||||
double err=0,ref=0; for(int i=0;i<n;i++){double d=got[i]-want[i];err+=d*d;ref+=(double)want[i]*want[i];}
|
||||
float r=(float)std::sqrt(err/(ref+1e-20));
|
||||
if(r>limit){std::fprintf(stderr,"relative RMS %.5f exceeds %.5f\n",r,limit);return 0;} return 1;
|
||||
}
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
int devices[COLI_CUDA_MAX_DEVICES], ndev = argc > 1 ? argc - 1 : 1;
|
||||
if (ndev > COLI_CUDA_MAX_DEVICES) return 2;
|
||||
@@ -55,8 +61,67 @@ int main(int argc, char **argv) {
|
||||
ColiCudaTensor *tf = nullptr;
|
||||
if (!coli_cuda_matmul(&tf, got, x, wf, nullptr, 0, 1, 4, 2, d0) || !close_enough(got, wantf, 2)) return 1;
|
||||
|
||||
const float eg[8] = {1,0,0,0, 0,1,0,0};
|
||||
const float eu[8] = {1,0,0,0, 0,1,0,0};
|
||||
const float ed[8] = {1,0, 0,1, 1,1, 1,-1};
|
||||
ColiCudaTensor *tg=nullptr,*tu=nullptr,*td=nullptr;
|
||||
if (!coli_cuda_tensor_upload(&tg,eg,nullptr,0,4,2,d0) ||
|
||||
!coli_cuda_tensor_upload(&tu,eu,nullptr,0,4,2,d0) ||
|
||||
!coli_cuda_tensor_upload(&td,ed,nullptr,0,2,4,d0)) return 1;
|
||||
float expert[8], want_expert[8];
|
||||
for(int s=0;s<2;s++){
|
||||
float a=x[s*4], b=x[s*4+1];
|
||||
a=(a/(1.0f+std::exp(-a)))*a; b=(b/(1.0f+std::exp(-b)))*b;
|
||||
want_expert[s*4]=a; want_expert[s*4+1]=b;
|
||||
want_expert[s*4+2]=a+b; want_expert[s*4+3]=a-b;
|
||||
}
|
||||
if (!coli_cuda_expert_mlp(tg,tu,td,expert,x,2) ||
|
||||
!close_enough(expert,want_expert,8)) return 1;
|
||||
ColiCudaTensor *gates[2]={tg,tg},*ups[2]={tu,tu},*downs[2]={td,td};
|
||||
int group_rows[2]={1,1}; float grouped[8];
|
||||
if (!coli_cuda_expert_group(gates,ups,downs,group_rows,2,grouped,x) ||
|
||||
!close_enough(grouped,want_expert,8)) return 1;
|
||||
|
||||
const float aw[16]={1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1};
|
||||
const float aq[4]={1,2,.5f,-.5f},al[12]={1,0,0,0, 0,1,0,0, 0,0,1,0};
|
||||
const float ar[6]={1,0, 0,1, 1,1};float actx[2],aref[2];
|
||||
ColiCudaTensor *at=nullptr;if(!coli_cuda_tensor_upload(&at,aw,nullptr,0,4,4,d0))return 1;
|
||||
float score[3];for(int t=0;t<3;t++)score[t]=aq[0]*al[t*4]+aq[1]*al[t*4+1]+aq[2]*ar[t*2]+aq[3]*ar[t*2+1];
|
||||
float mx=score[0],z=0;for(int t=1;t<3;t++)mx=score[t]>mx?score[t]:mx;
|
||||
for(int t=0;t<3;t++){score[t]=std::exp(score[t]-mx);z+=score[t];}for(int t=0;t<3;t++)score[t]/=z;
|
||||
for(int v=0;v<2;v++){aref[v]=0;for(int t=0;t<3;t++)aref[v]+=score[t]*al[t*4+2+v];}
|
||||
if(!coli_cuda_attention_absorb(at,actx,aq,al,ar,1,2,2,2,4,3,1.f)||
|
||||
!close_enough(actx,aref,2))return 1;
|
||||
coli_cuda_tensor_free(at);
|
||||
|
||||
/* Native s4 WMMA path: compare the quantized-activation result against the
|
||||
existing FP32-activation/s4-weight grouped implementation. */
|
||||
uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], tensor4[64];
|
||||
for(int i=0;i<(int)sizeof(w4);i++){
|
||||
int lo=((i%15)-7)&15,hi=(((i*3)%15)-7)&15;
|
||||
w4[i]=(uint8_t)(lo|(hi<<4));
|
||||
}
|
||||
for(int i=0;i<32;i++)ws4[i]=0.01f+(i%5)*0.002f;
|
||||
for(int i=0;i<64;i++)gx4[i]=std::sin((float)(i+1)*0.17f)*2.f;
|
||||
ColiCudaTensor *g4=nullptr,*u4=nullptr,*d4=nullptr;
|
||||
if(!coli_cuda_tensor_upload(&g4,w4,ws4,2,32,32,d0)||
|
||||
!coli_cuda_tensor_upload(&u4,w4,ws4,2,32,32,d0)||
|
||||
!coli_cuda_tensor_upload(&d4,w4,ws4,2,32,32,d0))return 1;
|
||||
ColiCudaTensor *gg4[2]={g4,g4},*ug4[2]={u4,u4},*dg4[2]={d4,d4};
|
||||
if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,scalar4,gx4))return 1;
|
||||
setenv("COLI_CUDA_TC_INT4","1",1);
|
||||
setenv("COLI_CUDA_TC_MIN_ROWS","1",1);
|
||||
if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,tensor4,gx4)||
|
||||
!relative_rms(tensor4,scalar4,64,0.30f))return 1;
|
||||
unsetenv("COLI_CUDA_TC_INT4");
|
||||
unsetenv("COLI_CUDA_TC_MIN_ROWS");
|
||||
coli_cuda_tensor_free(g4);coli_cuda_tensor_free(u4);coli_cuda_tensor_free(d4);
|
||||
uint64_t group_calls=0,group_experts=0,group_total_rows=0;
|
||||
coli_cuda_group_stats(&group_calls,&group_experts,&group_total_rows,nullptr,nullptr,nullptr);
|
||||
if(group_calls!=3||group_experts!=6||group_total_rows!=6) return 1;
|
||||
|
||||
coli_cuda_stats(-1, &count, &bytes);
|
||||
if (count != 4 || bytes != 70) {
|
||||
if (count != 7 || bytes != 166) {
|
||||
std::fprintf(stderr, "unexpected CUDA stats: %zu tensors, %zu bytes\n", count, bytes);
|
||||
return 1;
|
||||
}
|
||||
@@ -64,15 +129,18 @@ int main(int argc, char **argv) {
|
||||
coli_cuda_tensor_device(t4) != d1 || coli_cuda_tensor_device(t2) != d1) return 1;
|
||||
coli_cuda_stats(d0, &count, &bytes);
|
||||
if (ndev > 1) {
|
||||
if (count != 2 || bytes != 48) return 1;
|
||||
if (count != 5 || bytes != 144) return 1;
|
||||
coli_cuda_stats(d1, &count, &bytes);
|
||||
if (count != 2 || bytes != 22) return 1;
|
||||
} else if (count != 4 || bytes != 70) return 1;
|
||||
} else if (count != 7 || bytes != 166) return 1;
|
||||
|
||||
coli_cuda_tensor_free(t8);
|
||||
coli_cuda_tensor_free(t4);
|
||||
coli_cuda_tensor_free(t2);
|
||||
coli_cuda_tensor_free(tf);
|
||||
coli_cuda_tensor_free(tg);
|
||||
coli_cuda_tensor_free(tu);
|
||||
coli_cuda_tensor_free(td);
|
||||
coli_cuda_stats(-1, &count, &bytes);
|
||||
if (count || bytes) return 1;
|
||||
coli_cuda_shutdown();
|
||||
|
||||
Executable
BIN
Binary file not shown.
@@ -0,0 +1,56 @@
|
||||
#include <assert.h>
|
||||
#include <stdio.h>
|
||||
|
||||
#include "../decode_batch.h"
|
||||
|
||||
static void test_rows_use_their_own_sequence_storage(void)
|
||||
{
|
||||
float sequence_a[4 * 3] = {0};
|
||||
float sequence_b[4 * 3] = {0};
|
||||
|
||||
float *a2 = coli_kv_row(sequence_a, 2, 3);
|
||||
float *b1 = coli_kv_row(sequence_b, 1, 3);
|
||||
a2[0] = 20.0f;
|
||||
b1[2] = 12.0f;
|
||||
|
||||
assert(a2 == &sequence_a[6]);
|
||||
assert(b1 == &sequence_b[3]);
|
||||
assert(sequence_a[6] == 20.0f);
|
||||
assert(sequence_b[5] == 12.0f);
|
||||
assert(sequence_a[5] == 0.0f);
|
||||
assert(sequence_b[6] == 0.0f);
|
||||
}
|
||||
|
||||
static void test_const_reader_selects_the_same_row(void)
|
||||
{
|
||||
float storage[5 * 7] = {0};
|
||||
const float *row = coli_kv_row(storage, 4, 7);
|
||||
|
||||
assert(row == &storage[28]);
|
||||
}
|
||||
|
||||
static void test_submit_header(void)
|
||||
{
|
||||
ColiSubmit sub;
|
||||
assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95", &sub));
|
||||
assert(sub.id == 42 && sub.slot == 3 && sub.bytes == 17);
|
||||
assert(sub.max_tokens == 64 && sub.temperature > .69f && sub.top_p > .94f);
|
||||
assert(!coli_submit_parse("SUBMIT 1 -1 2 3 0.7 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 0 0.7 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 3 4 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 0 0 2 3 1 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 3 nan 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 3 1 inf", &sub));
|
||||
assert(coli_submit_parse("SUBMIT 1 0 16777216 3 1 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 16777217 3 1 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 3 1 1 trailing", &sub));
|
||||
}
|
||||
|
||||
int main(void)
|
||||
{
|
||||
test_rows_use_their_own_sequence_storage();
|
||||
test_const_reader_selects_the_same_row();
|
||||
test_submit_header();
|
||||
puts("decode batch helper tests: ok");
|
||||
return 0;
|
||||
}
|
||||
@@ -142,7 +142,7 @@ class DoctorTest(unittest.TestCase):
|
||||
output = format_doctor(self.report())
|
||||
|
||||
self.assertIn("model.path", output)
|
||||
self.assertIn("disk backing store", output)
|
||||
self.assertIn("disk 0.0 GB cold experts", output)
|
||||
self.assertTrue(output.endswith("result ok"))
|
||||
|
||||
def test_cli_json_is_machine_readable_without_loading_model(self):
|
||||
|
||||
Executable
BIN
Binary file not shown.
Executable
BIN
Binary file not shown.
@@ -1,19 +1,24 @@
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import socket
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
from urllib.error import HTTPError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler,
|
||||
generation_options, read_engine_turn, render_chat, serve)
|
||||
READY, Engine, generation_options, read_engine_turn, render_chat,
|
||||
serve)
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
|
||||
cancelled=None):
|
||||
self.calls.append((prompt, maximum, temperature, top_p, cache_slot))
|
||||
on_text("Hé")
|
||||
on_text("llo")
|
||||
@@ -26,10 +31,12 @@ class BlockingEngine(FakeEngine):
|
||||
self.entered = threading.Event()
|
||||
self.release = threading.Event()
|
||||
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
|
||||
cancelled=None):
|
||||
self.entered.set()
|
||||
self.release.wait(2)
|
||||
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot)
|
||||
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot,
|
||||
cancelled)
|
||||
|
||||
|
||||
class TemplateTest(unittest.TestCase):
|
||||
@@ -65,6 +72,10 @@ class TemplateTest(unittest.TestCase):
|
||||
(4, 0.0, 1.0))
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"max_tokens": 9}, 8)
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"temperature": math.nan}, 8)
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"top_p": math.inf}, 8)
|
||||
self.assertEqual(generation_options({"temperature": None, "top_p": None}, 8),
|
||||
(8, 0.7, 0.9))
|
||||
|
||||
@@ -82,8 +93,27 @@ class ProtocolTest(unittest.TestCase):
|
||||
with self.assertRaisesRegex(ValueError, "kv_slots"):
|
||||
serve("/missing", kv_slots=0)
|
||||
|
||||
def test_occupied_port_fails_before_engine_start(self):
|
||||
listener = socket.socket()
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen()
|
||||
try:
|
||||
with patch("openai_server.subprocess.Popen") as popen:
|
||||
with self.assertRaises(OSError):
|
||||
serve("/missing", port=listener.getsockname()[1])
|
||||
popen.assert_not_called()
|
||||
finally:
|
||||
listener.close()
|
||||
|
||||
|
||||
class SchedulerTest(unittest.TestCase):
|
||||
def test_admits_up_to_capacity_without_serializing(self):
|
||||
scheduler = GenerationScheduler(max_queue=0, queue_timeout=1, capacity=2)
|
||||
with scheduler.admit() as first:
|
||||
with scheduler.admit() as second:
|
||||
self.assertEqual({first[1], second[1]}, {0, 1})
|
||||
self.assertEqual(scheduler.snapshot()["active"], 2)
|
||||
|
||||
def test_rejects_when_waiting_queue_is_full(self):
|
||||
scheduler = GenerationScheduler(max_queue=0, queue_timeout=1)
|
||||
with scheduler.admit():
|
||||
@@ -160,6 +190,213 @@ class SchedulerTest(unittest.TestCase):
|
||||
self.assertEqual(errors, ["scheduler_closed"])
|
||||
|
||||
|
||||
class BlockingStream:
|
||||
def __init__(self, initial=b""):
|
||||
self.buffer = bytearray(initial)
|
||||
self.closed = False
|
||||
self.condition = threading.Condition()
|
||||
|
||||
def feed(self, data):
|
||||
with self.condition:
|
||||
self.buffer.extend(data)
|
||||
self.condition.notify_all()
|
||||
|
||||
def read(self, size=1):
|
||||
with self.condition:
|
||||
while len(self.buffer) < size and not self.closed:
|
||||
self.condition.wait()
|
||||
if not self.buffer and self.closed:
|
||||
return b""
|
||||
size = min(size, len(self.buffer))
|
||||
data = bytes(self.buffer[:size])
|
||||
del self.buffer[:size]
|
||||
return data
|
||||
|
||||
def readline(self):
|
||||
with self.condition:
|
||||
while b"\n" not in self.buffer and not self.closed:
|
||||
self.condition.wait()
|
||||
if not self.buffer and self.closed:
|
||||
return b""
|
||||
end = self.buffer.find(b"\n")
|
||||
size = len(self.buffer) if end < 0 else end + 1
|
||||
data = bytes(self.buffer[:size])
|
||||
del self.buffer[:size]
|
||||
return data
|
||||
|
||||
def close(self):
|
||||
with self.condition:
|
||||
self.closed = True
|
||||
self.condition.notify_all()
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, on_write):
|
||||
self.stdout = BlockingStream(READY + b"STAT 0 0 0 0\n")
|
||||
self.stdin = self
|
||||
self.on_write = on_write
|
||||
self.writes = []
|
||||
self.returncode = None
|
||||
|
||||
def write(self, data):
|
||||
self.writes.append(data)
|
||||
self.on_write(self, data)
|
||||
return len(data)
|
||||
|
||||
def flush(self):
|
||||
pass
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def terminate(self):
|
||||
self.returncode = 0
|
||||
self.stdout.close()
|
||||
|
||||
def wait(self, timeout=None):
|
||||
return self.returncode
|
||||
|
||||
def kill(self):
|
||||
self.terminate()
|
||||
|
||||
|
||||
class DispatcherTest(unittest.TestCase):
|
||||
def test_dispatches_interleaved_requests_by_id(self):
|
||||
submitted = []
|
||||
|
||||
def respond(process, frame):
|
||||
fields = frame.split(b"\n", 1)[0].split()
|
||||
self.assertEqual(fields[0], b"SUBMIT")
|
||||
submitted.append(fields[1])
|
||||
if len(submitted) == 2:
|
||||
first, second = submitted
|
||||
process.stdout.feed(b"DATA " + second + b" 3\nB-2\n")
|
||||
process.stdout.feed(b"DATA " + first + b" 3\nA-1\n")
|
||||
process.stdout.feed(b"DONE " + second + b" STAT 1 2.5 0 1.0 4 0\n")
|
||||
process.stdout.feed(b"DATA " + first + b" 3\nA-2\n")
|
||||
process.stdout.feed(b"DONE " + first + b" STAT 2 3.5 0 1.0 5 1\n")
|
||||
|
||||
process = FakeProcess(respond)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model", kv_slots=2)
|
||||
results = {}
|
||||
|
||||
def generate(name, prompt, slot):
|
||||
chunks = []
|
||||
stats = engine.generate(prompt, 8, 0.7, 0.9, chunks.append, slot)
|
||||
results[name] = ("".join(chunks), stats)
|
||||
|
||||
threads = [threading.Thread(target=generate, args=("a", "alpha", 0)),
|
||||
threading.Thread(target=generate, args=("b", "beta", 1))]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(timeout=2)
|
||||
self.assertFalse(thread.is_alive())
|
||||
engine.close()
|
||||
|
||||
self.assertEqual(results["a"][0], "A-1A-2")
|
||||
self.assertTrue(results["a"][1]["length_limited"])
|
||||
self.assertEqual(results["b"][0], "B-2")
|
||||
headers = [frame.split(b"\n", 1)[0].split() for frame in process.writes]
|
||||
self.assertEqual({int(header[2]) for header in headers}, {0, 1})
|
||||
self.assertEqual({header[3] for header in headers}, {b"4", b"5"})
|
||||
|
||||
def test_routes_engine_error_to_request(self):
|
||||
def respond(process, frame):
|
||||
request_id = frame.split()[1]
|
||||
process.stdout.feed(b"ERROR " + request_id + b" slot is busy\n")
|
||||
|
||||
process = FakeProcess(respond)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model")
|
||||
with self.assertRaisesRegex(RuntimeError, "slot is busy"):
|
||||
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
|
||||
engine.close()
|
||||
|
||||
def test_close_wakes_pending_generation_and_is_idempotent(self):
|
||||
process = FakeProcess(lambda _process, _frame: None)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model")
|
||||
errors = []
|
||||
|
||||
def generate():
|
||||
try:
|
||||
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
|
||||
except RuntimeError as error:
|
||||
errors.append(str(error))
|
||||
|
||||
thread = threading.Thread(target=generate)
|
||||
thread.start()
|
||||
for _ in range(100):
|
||||
with engine.pending_lock:
|
||||
if engine.pending:
|
||||
break
|
||||
threading.Event().wait(0.01)
|
||||
engine.close()
|
||||
engine.close()
|
||||
thread.join(timeout=2)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(errors, ["colibri engine is shutting down"])
|
||||
self.assertFalse(engine.dispatcher.is_alive())
|
||||
with engine.pending_lock:
|
||||
self.assertFalse(engine.pending)
|
||||
with self.assertRaisesRegex(RuntimeError, "shutting down"):
|
||||
engine.generate("again", 4, 0.7, 0.9, lambda _: None)
|
||||
|
||||
def test_protocol_corruption_fails_request_and_stops_dispatcher(self):
|
||||
def respond(process, frame):
|
||||
request_id = frame.split()[1]
|
||||
process.stdout.feed(b"DATA " + request_id + b" -1\n")
|
||||
|
||||
process = FakeProcess(respond)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model")
|
||||
with self.assertRaisesRegex(RuntimeError, "DATA size"):
|
||||
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
|
||||
with self.assertRaisesRegex(RuntimeError, "dispatcher stopped"):
|
||||
engine.generate("again", 4, 0.7, 0.9, lambda _: None)
|
||||
engine.close()
|
||||
|
||||
def test_decodes_utf8_split_across_data_frames(self):
|
||||
def respond(process, frame):
|
||||
request_id = frame.split()[1]
|
||||
process.stdout.feed(b"DATA " + request_id + b" 1\n\xc3\n")
|
||||
process.stdout.feed(b"DATA " + request_id + b" 1\n\xa9\n")
|
||||
process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n")
|
||||
|
||||
process = FakeProcess(respond)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model")
|
||||
chunks = []
|
||||
engine.generate("hello", 4, 0.7, 0.9, chunks.append)
|
||||
engine.close()
|
||||
self.assertEqual(chunks, ["é"])
|
||||
|
||||
def test_cancels_generation_after_consumer_disconnects(self):
|
||||
request_id = None
|
||||
|
||||
def respond(process, frame):
|
||||
nonlocal request_id
|
||||
fields = frame.split()
|
||||
if fields[0] == b"SUBMIT":
|
||||
request_id = fields[1]
|
||||
process.stdout.feed(b"DATA " + request_id + b" 1\nx\n")
|
||||
elif fields[0] == b"CANCEL":
|
||||
self.assertEqual(fields[1], request_id)
|
||||
process.stdout.feed(b"ERROR " + request_id + b" CANCELLED\n")
|
||||
|
||||
process = FakeProcess(respond)
|
||||
with patch("openai_server.subprocess.Popen", return_value=process):
|
||||
engine = Engine("glm", "model")
|
||||
output = []
|
||||
with self.assertRaises(ClientCancelled):
|
||||
engine.generate("hello", 8, 0.7, 0.9, output.append, cancelled=lambda: True)
|
||||
engine.close()
|
||||
self.assertEqual(output, ["x"])
|
||||
self.assertEqual(process.writes[-1].split(), [b"CANCEL", request_id])
|
||||
|
||||
|
||||
class HTTPTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
|
||||
@@ -57,11 +57,16 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
gpus = [{"index": 0, "name": "test-gpu", "total_bytes": 12 * GB,
|
||||
"free_bytes": 10 * GB}]
|
||||
plan = build_plan(self.model, ram_gb=16, context=32, vram_gb=20,
|
||||
available_memory=32 * GB, available_disk=100 * GB, gpus=gpus)
|
||||
self.assertEqual(plan["version"], 1)
|
||||
available_memory=32 * GB, available_disk=100 * GB, gpus=gpus,
|
||||
physical_cpus=24)
|
||||
self.assertEqual(plan["version"], 2)
|
||||
self.assertEqual(plan["policy"]["name"], "quality")
|
||||
self.assertEqual(plan["cpu"]["physical_cores"], 24)
|
||||
self.assertTrue(plan["policy"]["preserve_quantization"])
|
||||
self.assertFalse(plan["tiers"]["vram"]["requires_host_backing"])
|
||||
self.assertEqual(plan["tiers"]["ram"]["budget_bytes"], 16 * GB)
|
||||
self.assertLessEqual(plan["tiers"]["vram"]["budget_bytes"], 8 * GB)
|
||||
self.assertIn("required RAM backing", plan["warnings"][0])
|
||||
self.assertIn("clamped", plan["warnings"][0])
|
||||
self.assertIn("0:test-gpu", format_plan(plan))
|
||||
|
||||
def test_filters_requested_devices(self):
|
||||
@@ -78,7 +83,7 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
"--gpu", "none", "--json",
|
||||
], text=True, capture_output=True, check=True)
|
||||
plan = json.loads(run.stdout)
|
||||
self.assertEqual(plan["version"], 1)
|
||||
self.assertEqual(plan["version"], 2)
|
||||
self.assertEqual(plan["model"]["expert_count"], 2)
|
||||
|
||||
def test_applies_plan_without_overriding_explicit_settings(self):
|
||||
@@ -93,8 +98,15 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
self.assertEqual(env["RAM_GB"], "12")
|
||||
self.assertEqual(env["COLI_CUDA"], "1")
|
||||
self.assertEqual(env["COLI_GPUS"], "1")
|
||||
self.assertEqual(env["OMP_NUM_THREADS"], str(plan["cpu"]["physical_cores"]))
|
||||
self.assertEqual(env["OMP_PROC_BIND"], "spread")
|
||||
self.assertEqual(env["OMP_PLACES"], "cores")
|
||||
self.assertEqual(env["PIN_GB"], env["CUDA_EXPERT_GB"])
|
||||
|
||||
explicit_threads = environment_for_plan(plan, {"OMP_NUM_THREADS": "7",
|
||||
"OMP_PROC_BIND": "close"})
|
||||
self.assertEqual(explicit_threads["OMP_NUM_THREADS"], "7")
|
||||
self.assertEqual(explicit_threads["OMP_PROC_BIND"], "close")
|
||||
def test_cpu_binary_does_not_apply_gpu_tier(self):
|
||||
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=[{"index": 0, "name": "a", "total_bytes": 8 * GB,
|
||||
@@ -106,6 +118,32 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
self.assertNotIn("COLI_GPU", disabled)
|
||||
self.assertNotIn("CUDA_EXPERT_GB", disabled)
|
||||
|
||||
def test_rejects_unknown_policy_and_marks_experimental_policy(self):
|
||||
with self.assertRaisesRegex(ValueError, "unknown policy"):
|
||||
build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=[], policy="fast-ish")
|
||||
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=[], policy="experimental-fast")
|
||||
self.assertFalse(plan["policy"]["quality_preserving"])
|
||||
self.assertFalse(plan["policy"]["preserve_router"])
|
||||
|
||||
def test_balanced_policy_enables_lossless_live_repin(self):
|
||||
plan = build_plan(self.model, available_memory=16 * GB, available_disk=1,
|
||||
gpus=[], policy="balanced")
|
||||
env = environment_for_plan(plan)
|
||||
self.assertEqual(env["COLI_POLICY"], "balanced")
|
||||
self.assertEqual(env["REPIN"], "64")
|
||||
explicit = environment_for_plan(plan, {"REPIN": "0"})
|
||||
self.assertEqual(explicit["REPIN"], "0")
|
||||
|
||||
def test_plan_explains_hot_warm_and_cold_placement(self):
|
||||
plan = build_plan(self.model, ram_gb=4, vram_gb=0,
|
||||
available_memory=4 * GB, available_disk=1, gpus=[])
|
||||
self.assertEqual([item["target"] for item in plan["decisions"]],
|
||||
["VRAM", "RAM", "Disk"])
|
||||
self.assertIn("quality-preserving yes", format_plan(plan))
|
||||
self.assertIn("expected_bottleneck", plan)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -17,6 +17,11 @@ int main(void){
|
||||
|
||||
tier_decay(heat,6);
|
||||
if(heat[0]!=10 || heat[1]!=1 || heat[4]!=15) return fail("heat decay");
|
||||
|
||||
uint32_t freq[5]={10,10,2,18,18}, last[5]={10,90,95,20,99};
|
||||
int live[2]={0,1};
|
||||
if(!tier_pick_lfru(freq,last,100,5,live,2,&slot,&eid,&gain)) return fail("LFRU promotion");
|
||||
if(slot!=0||eid!=4) return fail("LFRU did not prefer recent ties");
|
||||
puts("tier tests: ok");
|
||||
return 0;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user