diff --git a/README.md b/README.md index 86aad79c0..626e92000 100644 --- a/README.md +++ b/README.md @@ -89,8 +89,40 @@ the right experts get. It works because routing has measurable structure (see the [expert atlas](https://github.com/JustVugg/colibri/issues/175)) — and structure is cacheable. -The engine is a single C file (`c/glm.c`) plus small headers. No BLAS, no Python -at runtime, no GPU required. +The engine is a single C file (`c/colibri.c`) plus small headers. No BLAS, no +Python at runtime, no GPU required. + +### Local cluster mode + +The coordinator keeps token generation, routing, and KV state local while +disk-backed expert workers execute routed FFNs on other Macs. A layer's routed +batch-union is sent as one persistent TCP request, so a token does not incur one +round trip per expert. + +Start the optional registration service: + +```bash +./coli cluster coordinator --host 0.0.0.0 --port 8765 +``` + +On each worker, with the same converted model available locally: + +```bash +./coli cluster worker --model /nvme/glm52_i4 --port 9100 \ + --coordinator http://COORDINATOR:8765 --advertise-host WORKER_IP +``` + +Run the coordinator with discovery, or provide `--cluster-workers +HOST:PORT,...` for a static setup: + +```bash +./coli serve --model /nvme/glm52_i4 \ + --cluster-coordinator http://127.0.0.1:8765 +``` + +The transport is disabled unless workers are configured, so the existing +single-machine path remains unchanged. Dense-layer sharding and browser/WebGPU +workers are separate follow-up seams. ## How it works diff --git a/c/cluster.py b/c/cluster.py new file mode 100644 index 000000000..d3ca5810b --- /dev/null +++ b/c/cluster.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Registration and discovery control plane for local expert workers.""" + +import argparse +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.request import Request, urlopen + + +PROTOCOL_VERSION = 1 + + +class ClusterRegistry: + def __init__(self, stale_after=30.0): + self.stale_after = float(stale_after) + self._nodes = {} + self._lock = threading.Lock() + + def register(self, node): + required = {"node_id", "host", "port", "role"} + missing = sorted(required - set(node)) + if missing: + raise ValueError("missing node fields: " + ", ".join(missing)) + port = int(node["port"]) + if not 1 <= port <= 65535: + raise ValueError("port must be between 1 and 65535") + role = str(node["role"]) + if role not in ("expert", "dense", "coordinator"): + raise ValueError("role must be expert, dense, or coordinator") + record = dict(node) + record.update(protocol_version=PROTOCOL_VERSION, port=port, last_seen=time.time()) + with self._lock: + self._nodes[str(node["node_id"])] = record + return record + + def heartbeat(self, node_id): + with self._lock: + node = self._nodes.get(str(node_id)) + if node is None: + raise KeyError(node_id) + node["last_seen"] = time.time() + return dict(node) + + def snapshot(self): + now = time.time() + with self._lock: + nodes = [dict(node) for node in self._nodes.values() + if now - node["last_seen"] <= self.stale_after] + nodes.sort(key=lambda node: (node["role"], node["node_id"])) + return {"protocol_version": PROTOCOL_VERSION, "nodes": nodes} + + def expert_endpoints(self): + return [f"{node['host']}:{node['port']}" + for node in self.snapshot()["nodes"] if node["role"] == "expert"] + + +class _Handler(BaseHTTPRequestHandler): + server_version = "colibri-cluster/1" + + def log_message(self, *_args): + return + + def _json(self, status, value): + payload = json.dumps(value, separators=(",", ":")).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def _body(self): + length = int(self.headers.get("Content-Length", "0")) + if length > 1 << 20: + raise ValueError("request body is too large") + return json.loads(self.rfile.read(length) or b"{}") + + def do_GET(self): # noqa: N802 - stdlib handler API + if self.path in ("/health", "/v1/cluster/topology"): + self._json(200, self.server.registry.snapshot()) + return + self._json(404, {"error": "not found"}) + + def do_POST(self): # noqa: N802 - stdlib handler API + try: + body = self._body() + if self.path == "/v1/cluster/register": + self._json(200, self.server.registry.register(body)) + elif self.path == "/v1/cluster/heartbeat": + self._json(200, self.server.registry.heartbeat(body["node_id"])) + else: + self._json(404, {"error": "not found"}) + except (KeyError, ValueError, TypeError, json.JSONDecodeError) as error: + self._json(400, {"error": str(error)}) + + +class ClusterServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, address, registry): + super().__init__(address, _Handler) + self.registry = registry + + +def serve(host="127.0.0.1", port=8765, stale_after=30.0): + server = ClusterServer((host, port), ClusterRegistry(stale_after)) + print(f"colibri cluster coordinator listening on http://{host}:{port}", flush=True) + try: + server.serve_forever() + finally: + server.server_close() + + +def register(coordinator, node): + request = Request(coordinator.rstrip("/") + "/v1/cluster/register", + data=json.dumps(node).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urlopen(request, timeout=5) as response: + return json.load(response) + + +def heartbeat(coordinator, node_id): + request = Request(coordinator.rstrip("/") + "/v1/cluster/heartbeat", + data=json.dumps({"node_id": node_id}).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urlopen(request, timeout=5) as response: + return json.load(response) + + +def discover_workers(coordinator): + with urlopen(coordinator.rstrip("/") + "/v1/cluster/topology", timeout=5) as response: + topology = json.load(response) + return [f"{node['host']}:{int(node['port'])}" + for node in topology.get("nodes", []) if node.get("role") == "expert"] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8765) + parser.add_argument("--stale-after", type=float, default=30.0) + args = parser.parse_args() + serve(args.host, args.port, args.stale_after) + + +if __name__ == "__main__": + main() diff --git a/c/coli b/c/coli index 5864e6898..8d7bc4cb6 100755 --- a/c/coli +++ b/c/coli @@ -149,6 +149,14 @@ def need_model(model): if not os.path.exists(GLM): sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build") +def need_worker_model(model): + if not os.path.isdir(model): + sys.exit(f"{C.yel}model not found:{C.r} {model}\n set COLI_MODEL or use --model") + if not os.path.exists(os.path.join(model,"config.json")): + sys.exit(f"{C.yel}config.json is missing from:{C.r} {model}") + if not os.path.exists(GLM): + sys.exit(f"{C.yel}engine is not built.{C.r} Run: coli build") + def cuda_binary(): if not os.path.exists(GLM): return False if sys.platform == "linux": @@ -211,6 +219,8 @@ def env_for(a): e.setdefault("PIPE", "1") e.setdefault("PILOT_REAL", "1") e["COLI_POLICY"]=a.policy + if getattr(a, "cluster_workers", None): e["CLUSTER_WORKERS"] = a.cluster_workers + if getattr(a, "cluster_coordinator", None): e["CLUSTER_COORDINATOR"] = a.cluster_coordinator if a.ram: e["RAM_GB"]=str(a.ram) if a.ngen: e["NGEN"]=str(a.ngen) if a.topp: e["TOPP"]=str(a.topp) @@ -828,14 +838,53 @@ def cmd_serve(a): with open(serve_pidfile(a.port),"w") as f: f.write(f"{os.getpid()} {a.model}\n") except OSError: pass from openai_server import serve + e=env_for(a) + if a.cluster_coordinator and not a.cluster_workers: + from cluster import discover_workers + try: + workers=discover_workers(a.cluster_coordinator) + except OSError as error: + sys.exit(f"{C.yel}cannot discover cluster workers:{C.r} {error}") + if not workers: + sys.exit(f"{C.yel}cluster coordinator has no live expert workers{C.r}") + e["CLUSTER_WORKERS"] = ",".join(workers) + print(f" {C.dim}[CLUSTER] discovered {len(workers)} expert worker(s){C.r}", file=sys.stderr) try: serve(a.model, a.host, a.port, a.model_id, a.api_key, - a.cap,a.ngen,GLM,env_for(a),a.cors_origin, + a.cap,a.ngen,GLM,e,a.cors_origin, a.max_queue,a.queue_timeout,a.kv_slots) finally: try: os.unlink(serve_pidfile(a.port)) except OSError: pass +def cmd_cluster_coordinator(a): + from cluster import serve + serve(a.host, a.port, a.stale_after) + +def cmd_cluster_worker(a): + need_worker_model(a.model) + coordinator=a.coordinator or os.environ.get("CLUSTER_COORDINATOR") + host=a.advertise_host or os.environ.get("CLUSTER_ADVERTISE_HOST", "127.0.0.1") + node_id=a.node_id or f"{host}:{a.port}" + stop_heartbeat=threading.Event() + if coordinator: + from cluster import heartbeat, register + register(coordinator, {"node_id":node_id, "host":host, "port":a.port, + "role":"expert", "layers":a.layers}) + def keep_registered(): + while not stop_heartbeat.wait(10): + try: heartbeat(coordinator, node_id) + except OSError: pass + threading.Thread(target=keep_registered, name="colibri-cluster-heartbeat", daemon=True).start() + e=env_for(a) + e.update({"EXPERT_WORKER":"1", "CLUSTER_WORKER_PORT":str(a.port), + "COLI_MMAP":os.environ.get("COLI_MMAP", "1")}) + print(f" {C.dim}[CLUSTER] expert worker · {host}:{a.port} · layers {a.layers}{C.r}") + try: + return subprocess.call([GLM,str(a.cap),str(a.ebits),str(a.dbits)],env=e) + finally: + stop_heartbeat.set() + def cmd_stop(a): """Shut down a running `coli serve` AND its engine — one command, no pkill. The engine re-execs itself for OMP tuning, so its process is named `exe`, @@ -963,6 +1012,10 @@ def main(): common.add_argument("--cap", type=int, default=8); common.add_argument("--ngen", type=int, default=1024) # rete di sicurezza: la fine vera la decidono gli stop token common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0) common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95) + common.add_argument("--cluster-workers", default=os.environ.get("CLUSTER_WORKERS"), + help="comma-separated expert workers, host:port,...") + common.add_argument("--cluster-coordinator", default=os.environ.get("CLUSTER_COORDINATOR"), + help="control-plane URL used to discover expert workers") ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibri — run GLM-5.2 locally") ap.add_argument("--version", action="version", version=f"colibri {_version}") sub=ap.add_subparsers(dest="cmd") @@ -988,6 +1041,17 @@ def main(): ps.add_argument("--max-queue",type=int,default=int(os.environ.get("COLI_MAX_QUEUE","8"))) ps.add_argument("--queue-timeout",type=float,default=float(os.environ.get("COLI_QUEUE_TIMEOUT","300"))) ps.add_argument("--kv-slots",type=int,default=int(os.environ.get("COLI_KV_SLOTS","1"))) + pcluster=sub.add_parser("cluster", help="run the local-cluster control plane or an expert worker") + cluster_sub=pcluster.add_subparsers(dest="cluster_cmd") + pcoord=cluster_sub.add_parser("coordinator", help="serve worker registration and discovery") + pcoord.add_argument("--host",default="127.0.0.1"); pcoord.add_argument("--port",type=int,default=8765) + pcoord.add_argument("--stale-after",type=float,default=30.0) + pworker=cluster_sub.add_parser("worker", parents=[common], help="serve disk-backed expert compute") + pworker.add_argument("--port",type=int,default=int(os.environ.get("CLUSTER_WORKER_PORT","9100"))) + pworker.add_argument("--coordinator",default=os.environ.get("CLUSTER_COORDINATOR")) + pworker.add_argument("--node-id",default=None); pworker.add_argument("--advertise-host",default=None) + pworker.add_argument("--layers",default="all",help="topology label, e.g. 0-37") + pworker.add_argument("--ebits",type=int,default=8); pworker.add_argument("--dbits",type=int,default=8) pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine") pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true") pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser") @@ -1015,6 +1079,9 @@ def main(): handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor, "run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench, "convert":cmd_convert,"web":cmd_web}.get(a.cmd) + if a.cmd=="cluster": + if a.cluster_cmd=="coordinator": handler=cmd_cluster_coordinator + elif a.cluster_cmd=="worker": handler=cmd_cluster_worker if handler: sys.exit(handler(a) or 0) banner(); print(__doc__) diff --git a/c/colibri.c b/c/colibri.c index bb24c8312..5aa8fef82 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -30,6 +30,10 @@ #include #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include /* select() serve-loop polling (#68); not on native MinGW */ +#include +#include +#include +#include #endif #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include @@ -75,6 +79,11 @@ static const float *g_pre_sh; /* routing precalcolata dalla GPU (Metal layer CB o device router CUDA, #431): * moe() la usa e salta la FASE A. NULL = router su CPU. */ static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; +#if !defined(_WIN32) +typedef struct { int fd; char host[128]; int port; } ClusterWorker; +static ClusterWorker g_cluster_workers[16]; +static int g_cluster_n; +#endif #ifdef __APPLE__ #include /* host_statistics64: MemAvailable di macOS */ #endif @@ -1670,6 +1679,175 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal, int de return rc; } +#if !defined(_WIN32) +/* Expert-worker protocol. Headers use network-order u32 values; activation + * bytes remain raw little-endian f32, matching the native engine ABI. One + * request contains the routed batch-union for a layer. */ +#define COLI_CLUSTER_MAGIC "COLIEX01" +#define COLI_CLUSTER_VERSION 1u +static int cluster_io(int fd, void *buf, size_t n, int write_mode){ + char *p=(char*)buf; + while(n){ + ssize_t r=write_mode?send(fd,p,n,0):recv(fd,p,n,MSG_WAITALL); + if(r<=0){ if(r<0&&errno==EINTR) continue; return -1; } + p+=r; n-=(size_t)r; + } + return 0; +} +static int cluster_u32(int fd, uint32_t *v, int write_mode){ + uint32_t x=write_mode?htonl(*v):0; + if(cluster_io(fd,write_mode?(void*)&x:(void*)v,sizeof(x),write_mode)) return -1; + if(!write_mode) *v=ntohl(*v); + return 0; +} +static int cluster_connect_one(const char *spec, ClusterWorker *out){ + char copy[256]; strncpy(copy,spec,sizeof(copy)-1); copy[sizeof(copy)-1]=0; + char *colon=strrchr(copy,':'); if(!colon||colon==copy||!colon[1]) return -1; + *colon=0; int port=atoi(colon+1); if(port<1||port>65535) return -1; + char portbuf[16]; snprintf(portbuf,sizeof(portbuf),"%d",port); + struct addrinfo hint={0},*ai=NULL; hint.ai_socktype=SOCK_STREAM; + if(getaddrinfo(copy,portbuf,&hint,&ai)!=0) return -1; + int fd=-1; + for(struct addrinfo *p=ai;p;p=p->ai_next){ + fd=socket(p->ai_family,p->ai_socktype,p->ai_protocol); + if(fd<0) continue; + if(connect(fd,p->ai_addr,p->ai_addrlen)==0) break; + close(fd); fd=-1; + } + freeaddrinfo(ai); if(fd<0) return -1; + out->fd=fd; strncpy(out->host,copy,sizeof(out->host)-1); out->host[sizeof(out->host)-1]=0; + out->port=port; return 0; +} +static void cluster_close_all(void){ + for(int i=0;i=0) close(g_cluster_workers[i].fd); + g_cluster_n=0; +} +static void cluster_init(void){ + const char *list=getenv("CLUSTER_WORKERS"); if(!list||!*list) return; + char *copy=strdup(list),*save=NULL; + for(char *tok=strtok_r(copy,",",&save);tok&&g_cluster_n<16;tok=strtok_r(NULL,",",&save)){ + while(*tok==' '||*tok=='\t') tok++; + if(cluster_connect_one(tok,&g_cluster_workers[g_cluster_n])) g_cluster_n++; + else fprintf(stderr,"[CLUSTER] cannot connect to expert worker %s\n",tok); + } + free(copy); + if(g_cluster_n<1){ fprintf(stderr,"[CLUSTER] no expert workers reachable\n"); exit(1); } + fprintf(stderr,"[CLUSTER] coordinator connected to %d expert worker(s)\n",g_cluster_n); +} +typedef struct { int eid,nr; int *rows; float *weights,*inputs; } ClusterItem; +static int cluster_item(const int *idxs,const float *ws,const int *keff,int K,int S, + int eid,ClusterItem *it,int D,const float *x){ + it->eid=eid; it->nr=0; + for(int s=0;snr++; break; } + if(!it->nr) return 0; + it->rows=malloc((size_t)it->nr*sizeof(int)); + it->weights=malloc((size_t)it->nr*sizeof(float)); + it->inputs=falloc((int64_t)it->nr*D); int r=0; + for(int s=0;srows[r]=s; it->weights[r]=ws[(int64_t)s*K+k]; + memcpy(it->inputs+(int64_t)r*D,x+(int64_t)s*D,(size_t)D*sizeof(float)); r++; break; + } + return 1; +} +static void cluster_item_free(ClusterItem *it){ free(it->rows); free(it->weights); free(it->inputs); memset(it,0,sizeof(*it)); } +static void cluster_moe_batch(Model *m,int layer,float *x,int S,float *out, + const int *idxs,const float *ws,const int *keff,int K, + const int *uniq,int base,int nb){ + int D=m->c.hidden; + for(int wi=0;wifd,(void*)COLI_CLUSTER_MAGIC,8,1)) goto fail; + v=COLI_CLUSTER_VERSION; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)layer; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)D; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)m->c.moe_inter; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)n; if(cluster_u32(w->fd,&v,1)) goto fail; + for(int j=0;jfd,&v,1)) goto fail; + v=(uint32_t)items[j].nr; if(cluster_u32(w->fd,&v,1)) goto fail; + if(cluster_io(w->fd,items[j].inputs,(size_t)items[j].nr*D*sizeof(float),1)) goto fail; + } + if(cluster_io(w->fd,magic,8,0)||memcmp(magic,COLI_CLUSTER_MAGIC,8)) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=COLI_CLUSTER_VERSION) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=0) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=(uint32_t)n) goto fail; + for(int j=0;jfd,&eid,0)||cluster_u32(w->fd,&nr,0) || + eid!=(uint32_t)items[j].eid || nr!=(uint32_t)items[j].nr) goto fail; + float *y=falloc((int64_t)nr*D); + if(cluster_io(w->fd,y,(size_t)nr*D*sizeof(float),0)){ free(y); goto fail; } + for(uint32_t r=0;rhost,w->port,layer); + exit(1); + } +} +typedef struct { int eid,nr; float *inputs; } ClusterRequestItem; +static int cluster_worker_run(const char *snap,int port,int ebits,int dbits){ + Model m; memset(&m,0,sizeof(m)); m.ebits=ebits; m.dbits=dbits; load_cfg(&m.c,snap); st_init(&m.S,snap); + int nr_layers=m.c.n_layers+1; ESlot *cache=calloc((size_t)nr_layers,sizeof(ESlot)); + for(int i=0;i=(uint32_t)nr_layers||n<1||n>64){ + close(cfd); cfd=-1; break; + } + ClusterRequestItem *items=calloc(n,sizeof(*items)); int bad=0; + for(uint32_t j=0;j=(uint32_t)m.c.n_experts||nr<1||nr>65536){bad=1;break;} + items[j].eid=(int)eid; items[j].nr=(int)nr; items[j].inputs=falloc((int64_t)nr*D); + if(cluster_io(cfd,items[j].inputs,(size_t)nr*D*sizeof(float),0)){bad=1;break;} + } + if(bad){ for(uint32_t j=0;jeid!=items[j].eid && expert_load(&m,(int)layer,items[j].eid,slot,1,0)){bad=1;break;} + int rows=items[j].nr; float *g=falloc((int64_t)rows*I),*u=falloc((int64_t)rows*I),*y=falloc((int64_t)rows*D); + expert_gate_up(g,u,items[j].inputs,&slot->g,&slot->u,rows); + for(int64_t z=0;z<(int64_t)rows*I;z++)g[z]=siluf(g[z])*u[z]; + if(slot->d.fmt==6) e8_rot_rows(g,rows,I); + matmul_qt(y,g,&slot->d,rows); + v=(uint32_t)items[j].eid; if(cluster_u32(cfd,&v,1)){bad=1;free(g);free(u);free(y);break;} + v=(uint32_t)rows; if(cluster_u32(cfd,&v,1)||cluster_io(cfd,y,(size_t)rows*D*sizeof(float),1)){bad=1;free(g);free(u);free(y);break;} + free(g);free(u);free(y); + } + for(uint32_t j=0;j=0)close(cfd); + } + close(fd); return 0; +} +#endif + #ifdef __linux__ /* io_uring expert batches. One owner prepares all reads for a block, submits * them in one syscall, and reaps CQEs on demand. The kernel, rather than a set @@ -3027,6 +3205,12 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int int shared_on_gpu=0; (void)shared_on_gpu; /* set by the Metal path when Phase E was fused */ for(int base=0;basepin[layer]; @@ -6475,6 +6659,17 @@ int main(int argc, char **argv){ int cap = argc>1?atoi(argv[1]):64; int ebits= argc>2?atoi(argv[2]):8; int dbits= argc>3?atoi(argv[3]):ebits; +#if !defined(_WIN32) + if(getenv("EXPERT_WORKER")){ + int port=getenv("CLUSTER_WORKER_PORT")?atoi(getenv("CLUSTER_WORKER_PORT")):9100; + if(port<1||port>65535){fprintf(stderr,"CLUSTER_WORKER_PORT must be 1..65535\n");return 2;} + return cluster_worker_run(snap,port,ebits,dbits); + } +#else + if(getenv("EXPERT_WORKER")){ + fprintf(stderr,"[CLUSTER] expert workers are not supported on Windows yet\n"); return 2; + } +#endif int kv_limit=(getenv("SERVE_BATCH")&&atoi(getenv("SERVE_BATCH")))?512:16; if(getenv("SERVE") && (kv_slot_count()<1 || kv_slot_count()>kv_limit)){ fprintf(stderr,"KV_SLOTS must be between 1 and %d\n",kv_limit); return 2; @@ -6533,6 +6728,12 @@ int main(int argc, char **argv){ #endif printf("== GLM C engine (glm_moe_dsa), cache=%d experts/layer | experts@%d-bit dense@%d-bit | idot: " IDOT_KERNEL " ==\n", cap, ebits, dbits); g_mem_avail_boot = mem_available_gb(); +#if !defined(_WIN32) + if(getenv("CLUSTER_WORKERS") && *getenv("CLUSTER_WORKERS")){ + cluster_init(); + atexit(cluster_close_all); + } +#endif Model m; double t0=now_s(); model_init(&m,snap,cap,ebits,dbits); if(!g_direct_heat_explicit){ /* COLI_DISKCLASS_WINDOW default, needs m.c (topk/n_layers) */ /* CURRENT-STATE CALIBRATION: the "8" multiplier (recency window ~= the last 8 diff --git a/c/tests/test_cluster.py b/c/tests/test_cluster.py new file mode 100644 index 000000000..0f84f8654 --- /dev/null +++ b/c/tests/test_cluster.py @@ -0,0 +1,48 @@ +import json +import sys +import threading +import unittest +from http.client import HTTPConnection +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from cluster import ClusterRegistry, ClusterServer, PROTOCOL_VERSION + + +class ClusterRegistryTests(unittest.TestCase): + def test_registers_and_discovers_expert_nodes(self): + registry = ClusterRegistry() + registry.register({"node_id": "mac-a", "host": "10.0.0.2", "port": 9100, + "role": "expert", "layers": "all"}) + registry.register({"node_id": "mac-b", "host": "10.0.0.3", "port": 9101, + "role": "dense", "layers": "38-75"}) + self.assertEqual(registry.expert_endpoints(), ["10.0.0.2:9100"]) + self.assertEqual(registry.snapshot()["protocol_version"], PROTOCOL_VERSION) + + def test_http_topology_registration_and_heartbeat(self): + server = ClusterServer(("127.0.0.1", 0), ClusterRegistry()) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + conn = HTTPConnection(*server.server_address) + body = json.dumps({"node_id": "mac-a", "host": "127.0.0.1", + "port": 9100, "role": "expert"}) + conn.request("POST", "/v1/cluster/register", body, + {"Content-Type": "application/json"}) + self.assertEqual(conn.getresponse().status, 200) + conn.request("POST", "/v1/cluster/heartbeat", json.dumps({"node_id": "mac-a"}), + {"Content-Type": "application/json"}) + self.assertEqual(conn.getresponse().status, 200) + conn.request("GET", "/v1/cluster/topology") + response = conn.getresponse() + self.assertEqual(response.status, 200) + self.assertEqual(len(json.loads(response.read())["nodes"]), 1) + conn.close() + finally: + server.shutdown() + server.server_close() + + +if __name__ == "__main__": + unittest.main()