diff --git a/README.md b/README.md index 2c7d62a..aaf7463 100644 --- a/README.md +++ b/README.md @@ -110,6 +110,22 @@ precision are the same whether an expert answered from VRAM or from disk. VRAM / RAM / NVMe three-tier expert residency

+### Dual-SSD: two copies of the model, twice the read bandwidth + +Decode is disk-bound on most machines, and expert reads are read-only — so if you have a **second SSD**, put a full copy of the model on it and let the engine stream from both drives at once: + +```bash +COLI_MODEL=/fast/glm52_i4 COLI_MODEL_MIRROR=/second/glm52_i4 ./coli chat +COLI_DISK_WEIGHTS=9,3 ... # optional: primary,mirror bandwidth ratio (else measured at startup) +``` + +Each expert is routed to one drive by a deterministic hash, weighted by the two drives' measured (or declared) bandwidth, so readahead/PILOT prefetch and the demand read always hit the same drive and nothing is cached twice. The aggregate bandwidth is the sum of both drives — a 9 GB/s + 3 GB/s pair reads experts ~33% faster than the fast drive alone, and the OMP-parallel pin/warmup load streams from both. Details worth knowing: + +- the mirror is **validated at startup** (per-file size + safetensors header must be byte-identical to the primary); divergent or missing files silently stay on the primary, so a **partial mirror is fine** — a smaller second SSD holding only some shards still helps; +- the mirror is **never written**: `.coli_usage`, `.coli_kv` and all sidecars stay on the primary; +- a read error on the mirror falls back to the primary (one warning, no crash), so unplugging the second drive mid-run degrades instead of killing the server; +- routing never changes tokens — both copies are byte-identical, and the per-run `MIRROR:` stats line shows GB served per drive. + The same engine spans the whole range: on a 25 GB laptop everything streams from disk (slow but correct); on a large host the entire expert set becomes resident (`CUDA_EXPERT_GB=auto PIN_GB=all`) and disk drops out of the decode path diff --git a/c/.gitignore b/c/.gitignore index d8e1f88..51695f9 100644 --- a/c/.gitignore +++ b/c/.gitignore @@ -3,3 +3,5 @@ glm_tiny/ olmoe_hf/ olmoe_i4/ .build-config +tests/test_st_mirror +tests/test_st_pread diff --git a/c/Makefile b/c/Makefile index 06fe176..4ad515f 100644 --- a/c/Makefile +++ b/c/Makefile @@ -171,7 +171,7 @@ else PYTHON ?= python3 endif CUDA_OBJ = -TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_pipe_block$(EXE) +TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_mirror$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_pipe_block$(EXE) ifneq (,$(LINUX)) TEST_BINS += tests/test_uring$(EXE) endif @@ -308,6 +308,9 @@ tests/test_st_pread$(EXE): tests/test_st_pread.c st.h json.h compat.h tests/test_st$(EXE): tests/test_st.c st.h json.h compat.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_st_mirror$(EXE): tests/test_st_mirror.c st.h json.h compat.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + tests/test_tier$(EXE): tests/test_tier.c tier.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) diff --git a/c/coli b/c/coli index b00a1cd..5cef2fc 100755 --- a/c/coli +++ b/c/coli @@ -15,6 +15,9 @@ Run GLM-5.2 (744B) locally on CPU with roughly 15-26 GB of RAM. Configuration through environment variables or flags (also valid after the subcommand): COLI_MODEL= model directory (default /home/vincenzo/glm52_i4) + COLI_MODEL_MIRROR= second copy of the model on another drive: expert reads + are split across both SSDs (COLI_DISK_WEIGHTS=9,3 sets the + primary,mirror bandwidth ratio; default: measured at startup) --ram N RAM budget in GB (automatically sizes the expert cache) --repin N adapt RAM/VRAM experts every N tokens --topp P adaptive expert top-p --topk N fixed top-k diff --git a/c/colibri.c b/c/colibri.c index a94f795..4a4b2e0 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -1263,6 +1263,57 @@ static void *map_of_fd(int fd){ return base; } +/* ==================== DUAL-SSD: two model copies, two drives ==================== + * COLI_MODEL_MIRROR= registers a SECOND (read-only) copy of the model on + * another drive; expert reads are split between the two according to + * COLI_DISK_WEIGHTS=, (relative bandwidth; without the env it + * is measured at startup with the engine's own access pattern). Cold decode is + * disk-bound (~11 GB/token): two NVMe drives reading in parallel add up. */ +static const char *g_mirror_dir=NULL; /* COLI_MODEL_MIRROR / SNAP_MIRROR */ +static int g_mirror=0; /* 1 = mirror active (at least one shard accepted) */ +static int g_mir_share=64; /* expert share routed to the mirror, out of 256 */ +static _Atomic int64_t g_mir_bytes[2]; /* bytes served per drive: [0] primary, [1] mirror */ +static _Atomic int64_t g_mir_nread[2]; + +/* replica of one expert: DETERMINISTIC hash of (layer,eid). Determinism is a + * requirement, not a style choice: the readahead/PILOT WILLNEED and the demand + * pread must hit the same fd/page-cache, and in buffered mode an expert must + * never be cached twice (one copy per drive). */ +static inline int expert_route(int layer,int eid){ + if(!g_mirror) return 0; + uint32_t h=(uint32_t)layer*2654435761u ^ (uint32_t)eid*0x9E3779B9u; + h^=h>>16; h*=0x45d9f3bu; h^=h>>16; + return (int)(h&255) < g_mir_share; +} + +/* buffered fd of the replica, falling back to the primary if the file is not mirrored */ +static inline int rep_bfd(shards *S,int fd,int rep){ + int r=st_fd_rep(S,fd,rep); return r<0?fd:r; +} + +static int pread_full(int fd, void *buf, int64_t n, int64_t off, const char *tag); + +/* pread on the chosen replica with fallback to the primary on error/short-read: + * an unreadable sector (or an unmount) of the mirror must never kill the process + * when the primary can serve the same bytes. Delegates to pread_full so the + * mirror path inherits the short-read/EINTR loop and honest reporting (#236). + * Accounts bytes per drive. Returns 0 = ok, -1 = real error/EOF (like pread_full). */ +static ssize_t mir_pread(shards *S,int fd,int rep,void *buf,int64_t n,int64_t off,const char *tag){ + int rfd = st_fd_rep(S,fd,rep); + int used = rep && rfd>=0; + if(rfd<0) rfd=fd; + int rc=pread_full(rfd,buf,n,off,tag); + if(rc && used){ + static _Atomic int warned; + if(!atomic_exchange(&warned,1)) + fprintf(stderr,"[MIRROR] read error on the mirror copy — falling back to the primary drive\n"); + used=0; rc=pread_full(fd,buf,n,off,tag); + } + if(!rc){ atomic_fetch_add_explicit(&g_mir_bytes[used],n,memory_order_relaxed); + atomic_fetch_add_explicit(&g_mir_nread[used],1,memory_order_relaxed); } + return rc; +} + /* carica un expert nello slot. Container pre-quantizzato: le 3 matrici sono contigue nel * file -> UNA pread coalescente da ~19 MB dentro `slab` (+ le scale in fslab); i QT sono * viste dentro lo slab (zero copie). Fallback per modelli non quantizzati (oracolo tiny). @@ -1358,10 +1409,12 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i else { __atomic_add_fetch(&m->ld_main,1,__ATOMIC_RELAXED); __atomic_add_fetch(&m->bytes_main,(uint64_t)tb,__ATOMIC_RELAXED); } } + int rep=expert_route(layer,eid); /* DUAL-SSD: this expert's replica */ + if(rep && st_fd_rep(&m->S,tw[0]->fd,1)<0) rep=0; /* shard not in the mirror (partial) */ if(g_mmap){ void *bw[3],*bq[3]; int okm=1; for(int k=0;k<3;k++){ - bw[k]=map_of_fd(tw[k]->fd); bq[k]=map_of_fd(tq[k]->fd); + bw[k]=map_of_fd(rep_bfd(&m->S,tw[k]->fd,rep)); bq[k]=map_of_fd(rep_bfd(&m->S,tq[k]->fd,rep)); if(!bw[k]||!bq[k]||((tw[k]->off)&3)||((tq[k]->off)&3)) okm=0; } if(okm){ @@ -1396,7 +1449,9 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i * the final resident set only, after GPU release has already nulled out the * pointers for anything that isn't genuinely RAM-tier. */ atomic_fetch_add_explicit(&g_prof_io,(int64_t)(n+nq),memory_order_relaxed); + atomic_fetch_add_explicit(&g_mir_bytes[rep],tw[k]->nbytes+tq[k]->nbytes,memory_order_relaxed); } + atomic_fetch_add_explicit(&g_mir_nread[rep],1,memory_order_relaxed); s->eid=eid; return 0; } } @@ -1475,7 +1530,7 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i int64_t pos[3]; int done=0, dc_direct=0; if(contig){ int64_t off0=tw[ord[0]]->off; - int dfd = g_direct ? st_direct_fd(&m->S, tw[ord[0]]->fd) : -1; + int dfd = g_direct ? st_direct_fd_rep(&m->S, tw[ord[0]]->fd, rep) : -1; if(dfd>=0){ /* O_DIRECT: offset/len allineati a 4K */ int64_t base=off0 & ~4095LL, need=(off0-base)+wtot; int64_t len=(need+4095)&~4095LL; @@ -1483,10 +1538,12 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i if(r>=need){ pos[ord[0]]=off0-base; pos[ord[1]]=pos[ord[0]]+tw[ord[0]]->nbytes; pos[ord[2]]=pos[ord[1]]+tw[ord[1]]->nbytes; done=1; dc_direct=1; + atomic_fetch_add_explicit(&g_mir_bytes[rep],(int64_t)r,memory_order_relaxed); + atomic_fetch_add_explicit(&g_mir_nread[rep],1,memory_order_relaxed); } } if(!done){ /* fallback bufferizzato */ - if(pread_full(tw[ord[0]]->fd, s->slab, wtot, off0, "pread expert")){ if(fatal) exit(1); + if(mir_pread(&m->S, tw[ord[0]]->fd, rep, s->slab, wtot, off0, "pread expert")){ if(fatal) exit(1); if(dc_on) dc_wall_exit(dc_cls,now_s()); /* pair the enter on the non-fatal unwind */ return -1; } pos[ord[0]]=0; pos[ord[1]]=tw[ord[0]]->nbytes; pos[ord[2]]=tw[ord[0]]->nbytes+tw[ord[1]]->nbytes; done=1; @@ -1495,14 +1552,14 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i if(!done){ /* non contigui: 3 pread bufferizzate */ int64_t o=0; for(int a=0;a<3;a++){ int k=ord[a]; - if(pread_full(tw[k]->fd, s->slab+o, tw[k]->nbytes, tw[k]->off, "pread expert")){ if(fatal) exit(1); + if(mir_pread(&m->S, tw[k]->fd, rep, s->slab+o, tw[k]->nbytes, tw[k]->off, "pread expert")){ if(fatal) exit(1); if(dc_on) dc_wall_exit(dc_cls,now_s()); /* pair the enter on the non-fatal unwind */ return -1; } pos[k]=o; o+=tw[k]->nbytes; } } float *fp[3]; int64_t fo=0; /* scale (piccole) */ for(int k=0;k<3;k++){ - if(pread_full(tq[k]->fd, (char*)(s->fslab+fo), tq[k]->nbytes, tq[k]->off, "pread qs")){ if(fatal) exit(1); + if(mir_pread(&m->S, tq[k]->fd, rep, (char*)(s->fslab+fo), tq[k]->nbytes, tq[k]->off, "pread qs")){ if(fatal) exit(1); if(dc_on) dc_wall_exit(dc_cls,now_s()); /* pair the enter on the non-fatal unwind */ return -1; } fp[k]=s->fslab+fo; fo+=tq[k]->nbytes/4; } @@ -1518,9 +1575,10 @@ static int expert_load_impl(Model *m, int layer, int eid, ESlot *s, int fatal, i atomic_fetch_add_explicit(&g_dc_direct_n[dc_cls],1,memory_order_relaxed); } if(g_drop){ /* scarta subito le pagine: evita che la page - * cache in pressione strangoli il throughput */ - posix_fadvise(tw[ord[0]]->fd, tw[ord[0]]->off, wtot, POSIX_FADV_DONTNEED); - for(int k=0;k<3;k++) posix_fadvise(tq[k]->fd, tq[k]->off, tq[k]->nbytes, POSIX_FADV_DONTNEED); + * cache in pressione strangoli il throughput. + * The drop targets the fd of the replica READ. */ + posix_fadvise(rep_bfd(&m->S,tw[ord[0]]->fd,rep), tw[ord[0]]->off, wtot, POSIX_FADV_DONTNEED); + for(int k=0;k<3;k++) posix_fadvise(rep_bfd(&m->S,tq[k]->fd,rep), tq[k]->off, tq[k]->nbytes, POSIX_FADV_DONTNEED); } QT *qt[3]={&s->g,&s->u,&s->d}; int OO[3]={I,I,D}, II[3]={D,D,I}; for(int k=0;k<3;k++){ @@ -1923,15 +1981,17 @@ static void expert_host_ensure(Model *m, int layer, ESlot *s){ * Le scale .qs restano SEMPRE bufferizzate (pread sul fd normale), quindi il loro * WILLNEED resta utile anche con DIRECT=1. fadvise e' solo consultivo: saltarlo non * cambia mai l'output (bit-identico), riduce solo I/O sprecato. + * The readahead targets the SAME replica that will serve the pread (expert_route + * is deterministic). * EN: under O_DIRECT the weights bypass the page cache, so their WILLNEED is wasted; * the .qs scales are always buffered, so keep theirs. Advisory hint -> output-preserving. */ static void expert_prefetch(Model *m, int layer, int eid){ - char nm[300]; + char nm[300]; int rep=expert_route(layer,eid); const char *suf[3]={"gate_proj.weight","up_proj.weight","down_proj.weight"}; for(int k=0;k<3;k++){ snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.%s",layer,eid,suf[k]); - if(!g_direct) st_prefetch(&m->S,nm); - char qs[320]; snprintf(qs,sizeof(qs),"%s.qs",nm); st_prefetch(&m->S,qs); + if(!g_direct) st_prefetch_rep(&m->S,nm,rep); + char qs[320]; snprintf(qs,sizeof(qs),"%s.qs",nm); st_prefetch_rep(&m->S,qs,rep); } } @@ -4415,6 +4475,14 @@ static void profile_print(Model *m, double elapsed){ (unsigned long long)m->cpu_expert_rows,m->t_egpu,m->t_route,m->t_p2p,(unsigned long long)m->n_p2p, elapsed-m->t_ewait-m->t_emm-m->t_attn-m->t_head-m->t_route-m->t_p2p>0? elapsed-m->t_ewait-m->t_emm-m->t_attn-m->t_head-m->t_route-m->t_p2p:0); + if(g_mirror){ + double b0=atomic_load_explicit(&g_mir_bytes[0],memory_order_relaxed)/1e9; + double b1=atomic_load_explicit(&g_mir_bytes[1],memory_order_relaxed)/1e9; + printf("MIRROR: primary %.2f GB (%lld reads) | mirror %.2f GB (%lld reads) — %.0f%% of expert bytes from the mirror\n", + b0,(long long)atomic_load_explicit(&g_mir_nread[0],memory_order_relaxed), + b1,(long long)atomic_load_explicit(&g_mir_nread[1],memory_order_relaxed), + b0+b1>0?100.0*b1/(b0+b1):0.0); + } #ifdef COLI_METAL if(g_metal_enabled){ uint64_t ok=0,fb=0,ex=0; double su=0,gp=0,sc=0; coli_metal_moe_counts(&ok,&fb,&ex); coli_metal_moe_times(&su,&gp,&sc); @@ -4558,6 +4626,8 @@ static void run_replay(Model *m, const int *full, int nfull, int np){ m->hits=m->miss=m->ereq=m->gpu_expert_calls=0; m->hit_pin=m->hit_ecache=0; profile_reset(m); ProfBase pb; prof_base(m,&pb); + atomic_store(&g_mir_bytes[0],0); atomic_store(&g_mir_bytes[1],0); + atomic_store(&g_mir_nread[0],0); atomic_store(&g_mir_nread[1],0); double t0=now_s(); int steps=0; for(int i=np-1;infd;i++){ + if(rep && S->mfds[i]<0) continue; + int64_t sz=lseek(rep?S->mfds[i]:S->fds[i],0,SEEK_END); + if(sz>bsz){ bsz=sz; big=i; } + } + const int64_t blk=19ll<<20; const int NB=8; + if(big<0 || bszmdfds[big] : S->dfds[big]; + int fd = dfd>=0? dfd : (rep? S->mfds[big] : S->fds[big]); + if(dfd<0) fprintf(stderr,"[MIRROR] no O_DIRECT on %s: the probe may read the page cache — " + "set COLI_DISK_WEIGHTS for an accurate split\n", rep?"the mirror":"the primary"); + double t0=now_s(); int64_t tot=0; + #pragma omp parallel for schedule(dynamic,1) reduction(+:tot) + for(int i=0;i0) tot+=r; + compat_aligned_free(buf); + } + } + double dt=now_s()-t0; + return (dt>0 && tot>0)? tot/1e9/dt : 0; +} + +/* DUAL-SSD setup: register the mirror copy, derive the read split from + * COLI_DISK_WEIGHTS=, or from a startup bandwidth probe. + * Runs after model_init (needs the shard index) and BEFORE any pin/autopin + * load, so the OMP-parallel pin warmup already streams from both drives. */ +static void mirror_setup(Model *m){ + if(!g_mirror_dir) return; + const char *snap=getenv("SNAP"); + if(snap && !strcmp(snap,g_mirror_dir)){ + fprintf(stderr,"[MIRROR] mirror dir equals the model dir — ignored\n"); return; + } + int nf=st_mirror_init(&m->S,g_mirror_dir); + if(nf<=0){ + fprintf(stderr,"[MIRROR] %s: no usable shard (missing or divergent copy) — " + "running on the primary drive only\n",g_mirror_dir); + return; + } + g_mirror=1; + double wp=0,wm=0; const char *w=getenv("COLI_DISK_WEIGHTS"); const char *how="COLI_DISK_WEIGHTS"; + if(w && (sscanf(w," %lf , %lf",&wp,&wm)!=2 || wp<=0 || wm<=0)){ + fprintf(stderr,"[MIRROR] invalid COLI_DISK_WEIGHTS '%s' (want e.g. 9,3) — probing instead\n",w); + wp=wm=0; + } + if(wp<=0||wm<=0){ + wp=mirror_probe_bw(&m->S,0); wm=mirror_probe_bw(&m->S,1); how="measured"; + if(wp>0&&wm>0) fprintf(stderr,"[MIRROR] probe: primary %.2f GB/s, mirror %.2f GB/s\n",wp,wm); + else { wp=wm=1; how="fallback 1:1 (probe failed)"; } + } + int share=(int)(256.0*wm/(wp+wm)+0.5); + g_mir_share = share<1?1 : share>255?255 : share; + fprintf(stderr,"[MIRROR] %s: %d/%d shards | read split primary %.0f%% / mirror %.0f%% (%s)\n", + g_mirror_dir,nf,m->S.nfd,100.0*(256-g_mir_share)/256,100.0*g_mir_share/256,how); +} + typedef struct { int l,e; uint32_t c; } PinRec; static int pin_rec_cmp(const void *a,const void *b){ const PinRec *x=a,*y=b; return x->cc?1:x->c>y->c?-1:0; @@ -6011,6 +6146,9 @@ int main(int argc, char **argv){ fprintf(stderr,"URING=1 is supported only on Linux\n"); return 2; #endif } + g_mirror_dir = getenv("COLI_MODEL_MIRROR"); /* DUAL-SSD: second model copy */ + if(!g_mirror_dir||!*g_mirror_dir) g_mirror_dir = getenv("SNAP_MIRROR"); + if(g_mirror_dir&&!*g_mirror_dir) g_mirror_dir = NULL; g_idot = getenv("IDOT")?atoi(getenv("IDOT")):1; /* 0 = kernel f32 esatti (A/B) */ g_spec_pin = getenv("SPEC_PIN")?atoi(getenv("SPEC_PIN")):1; /* #163: 0 = gate S-dipendenti storici / legacy S-dependent gates */ if(getenv("ROUTE_TRACE")&&*getenv("ROUTE_TRACE")){ @@ -6147,6 +6285,9 @@ int main(int argc, char **argv){ " Keep it on a native Linux fs (ext4/xfs/zfs) for memory efficiency and speed.\n", snap); } #endif + /* DUAL-SSD: register the mirror copy BEFORE any pin/autopin load, so the + * OMP-parallel pin warmup already streams from both drives. */ + mirror_setup(&m); /* HOT-STORE: PIN= [PIN_GB=g] -> top expert per frequenza fissi in RAM. * Va PRIMA di cap_for_ram: i pinnati contano nel residente. */ if(getenv("PIN")){ diff --git a/c/st.h b/c/st.h index a2c2c3f..8ed0388 100644 --- a/c/st.h +++ b/c/st.h @@ -40,6 +40,9 @@ typedef struct { int dfds[512]; /* gemelli O_DIRECT (aperti pigramente): -2 = non ancora provato */ char *paths[512]; int nfd; + int mfds[512]; /* MIRROR: fds of the second model copy (dual-SSD), -1 = absent */ + int mdfds[512]; /* O_DIRECT twins of the second copy, -1 = absent */ + int nmirror; /* files accepted into the mirror (0 = mirror inactive) */ int *hidx; /* hash map nome->indice (open addressing): con ~120k tensori * (GLM: 256 expert x 78 layer x 3 x 2) la scansione lineare * costava decine di secondi/token (misurato sul primo run reale) */ @@ -99,10 +102,80 @@ static int st_open_fd(shards *S, const char *path) { /* fd gemello O_DIRECT dello stesso file (bypassa la page cache: il buffered read su * ext4-in-VHDX si strozza a ~0.8 GB/s, O_DIRECT arriva a 2.3+; misurato). -1 se non disponibile. */ -static int st_direct_fd(shards *S, int fd) { - for (int i = 0; i < S->nfd; i++) if (S->fds[i] == fd) return S->dfds[i]; +static int st_fidx(shards *S, int fd) { + for (int i = 0; i < S->nfd; i++) if (S->fds[i] == fd) return i; return -1; } +static int st_direct_fd(shards *S, int fd) { + int i = st_fidx(S, fd); return i < 0 ? -1 : S->dfds[i]; +} + +/* ---- MIRROR (dual-SSD): second read-only copy of the model on another drive ---- + * st_fd_rep/st_direct_fd_rep: fd of replica `rep` (0 = primary, 1 = mirror) for + * the SAME file identified by its primary fd. -1 if that replica is absent. */ +static int st_fd_rep(shards *S, int fd, int rep) { + if (!rep) return fd; + if (!S->nmirror) return -1; + int i = st_fidx(S, fd); return i < 0 ? -1 : S->mfds[i]; +} +static int st_direct_fd_rep(shards *S, int fd, int rep) { + if (!rep) return st_direct_fd(S, fd); + if (!S->nmirror) return -1; + int i = st_fidx(S, fd); return i < 0 ? -1 : S->mdfds[i]; +} + +/* Registers / as a read replica of every already-indexed shard. + * A file is accepted ONLY if its size and safetensors header are byte-identical + * to the primary: the data_offsets then match by construction, so every pread + * is valid on either copy. Missing or divergent files simply stay on the + * primary (the mirror may be partial, e.g. a smaller SSD holding only the + * expert shards). Returns the number of accepted files. The mirror is NEVER + * written to: .coli_usage/.coli_kv keep deriving from the primary alone. */ +static int st_mirror_init(shards *S, const char *dir) { + if (S->nmirror) for (int i = 0; i < S->nfd; i++) { /* re-init: drop the old replica */ + if (S->mfds[i] >= 0) close(S->mfds[i]); + if (S->mdfds[i] >= 0) close(S->mdfds[i]); + } + for (int i = 0; i < ST_MAX_SHARDS; i++) { S->mfds[i] = -1; S->mdfds[i] = -1; } + S->nmirror = 0; + for (int i = 0; i < S->nfd; i++) { + const char *base = strrchr(S->paths[i], '/'); +#ifdef _WIN32 + const char *b2 = strrchr(S->paths[i], '\\'); + if (b2 && (!base || b2 > base)) base = b2; +#endif + base = base ? base + 1 : S->paths[i]; + char mp[2048]; snprintf(mp, sizeof(mp), "%s/%s", dir, base); + int mfd = open(mp, COMPAT_O_RDONLY); + if (mfd < 0) continue; /* partial mirror: this shard stays on the primary */ + int64_t sza = lseek(S->fds[i], 0, SEEK_END), szb = lseek(mfd, 0, SEEK_END); + if (sza != szb) { + fprintf(stderr, "[MIRROR] %s: size differs from the primary copy — file skipped\n", mp); + close(mfd); continue; + } + uint64_t ha = 0, hb = 0; int ok = 1; /* identical header => identical data_offsets */ + if (pread(S->fds[i], &ha, 8, 0) != 8 || pread(mfd, &hb, 8, 0) != 8 || + ha != hb || ha == 0 || ha > (uint64_t)256 << 20 || (int64_t)(8 + ha) > sza) ok = 0; + if (ok) { + char *ba = malloc(ha), *bb = malloc(ha); + if (!ba || !bb || pread(S->fds[i], ba, ha, 8) != (ssize_t)ha || + pread(mfd, bb, ha, 8) != (ssize_t)ha || memcmp(ba, bb, ha)) ok = 0; + free(ba); free(bb); + } + if (!ok) { + fprintf(stderr, "[MIRROR] %s: header differs from the primary copy — file skipped\n", mp); + close(mfd); continue; + } + S->mfds[i] = mfd; +#ifdef O_DIRECT + S->mdfds[i] = open(mp, COMPAT_O_RDONLY | O_DIRECT); +#elif defined(__APPLE__) + S->mdfds[i] = compat_open_direct(mp); +#endif + S->nmirror++; + } + return S->nmirror; +} /* indicizza tutti i model-*.safetensors in snap_dir */ /* pread completo: chunk-loop (una singola pread si ferma a ~2^31 byte su Linux @@ -256,6 +329,16 @@ static void st_prefetch(shards *S, const char *name) { if (t) posix_fadvise(t->fd, t->off, t->nbytes, POSIX_FADV_WILLNEED); } +/* like st_prefetch, but on replica `rep`'s drive: the WILLNEED must warm the + * page cache of the SAME fd the later demand pread will hit. */ +static void st_prefetch_rep(shards *S, const char *name, int rep) { + st_tensor *t = st_find(S, name); + if (!t) return; + int fd = st_fd_rep(S, t->fd, rep); + if (fd < 0) fd = t->fd; + posix_fadvise(fd, t->off, t->nbytes, POSIX_FADV_WILLNEED); +} + /* legge un tensore in un buffer float32 fornito dal chiamante (numel float). * drop=1 -> consiglia al kernel di scartare le pagine (per gli expert in streaming). */ static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) { diff --git a/c/tests/test_st_mirror.c b/c/tests/test_st_mirror.c new file mode 100644 index 0000000..769c99d --- /dev/null +++ b/c/tests/test_st_mirror.c @@ -0,0 +1,109 @@ +/* Dual-SSD mirror (st_mirror_init & friends): a second read-only model copy is + * accepted only when byte-identical in size and safetensors header, reads on + * either replica return the same bytes, and divergent/missing copies degrade + * to the primary instead of being trusted. Fixture dirs are created in the + * current working directory and removed on exit. */ +#include +#include +#include + +#include "../st.h" + +#ifdef _WIN32 +#include +#define MKDIR(p) _mkdir(p) +#else +#include +#define MKDIR(p) mkdir(p, 0777) +#endif + +#define CHECK(condition) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: check failed: %s\n", __FILE__, __LINE__, #condition); \ + return 1; \ + } \ +} while (0) + +#define DIR_A "tmp_mirror_a" /* primary */ +#define DIR_B "tmp_mirror_b" /* identical copy */ +#define DIR_C "tmp_mirror_c" /* same size, divergent header */ +#define DIR_D "tmp_mirror_d" /* different size */ +#define DIR_E "tmp_mirror_e" /* empty (missing file) */ + +/* one-tensor safetensors file; flip lets us corrupt one header byte and + * pad lets us grow the payload, both without changing anything else */ +static int write_model(const char *dir, int flip, int pad) { + char hdr[128], path[256]; + int hl = snprintf(hdr, sizeof(hdr), + "{\"t0\":{\"dtype\":\"F32\",\"shape\":[8],\"data_offsets\":[0,32]}}"); + if (flip) hdr[hl - 2] = ' '; /* inside the JSON header, same length */ + snprintf(path, sizeof(path), "%s/model.safetensors", dir); + FILE *f = fopen(path, "wb"); + if (!f) { perror(path); return -1; } + uint64_t n = (uint64_t)hl; + fwrite(&n, 8, 1, f); + fwrite(hdr, 1, (size_t)hl, f); + float data[8] = {1, -2, 3.5f, 0, 42, -0.5f, 7, 8}; + fwrite(data, 4, 8, f); + for (int i = 0; i < pad; i++) fputc(0, f); + fclose(f); + return 0; +} + +static void cleanup(void) { + const char *dirs[] = {DIR_A, DIR_B, DIR_C, DIR_D, DIR_E}; + for (int i = 0; i < 5; i++) { + char path[256]; + snprintf(path, sizeof(path), "%s/model.safetensors", dirs[i]); + remove(path); + remove(dirs[i]); + } +} + +int main(void) { + cleanup(); + CHECK(MKDIR(DIR_A) == 0 && MKDIR(DIR_B) == 0 && MKDIR(DIR_C) == 0 && + MKDIR(DIR_D) == 0 && MKDIR(DIR_E) == 0); + CHECK(write_model(DIR_A, 0, 0) == 0); + CHECK(write_model(DIR_B, 0, 0) == 0); + CHECK(write_model(DIR_C, 1, 0) == 0); + CHECK(write_model(DIR_D, 0, 64) == 0); + + shards S; + st_init(&S, DIR_A); + CHECK(S.n == 1 && S.nfd == 1); + st_tensor *t = st_find(&S, "t0"); + CHECK(t != NULL); + + /* without a mirror: replica 0 is the identity, replica 1 is absent */ + CHECK(st_fd_rep(&S, t->fd, 0) == t->fd); + CHECK(st_fd_rep(&S, t->fd, 1) == -1); + + /* identical copy: accepted, and both replicas serve the same bytes */ + CHECK(st_mirror_init(&S, DIR_B) == 1); + int mfd = st_fd_rep(&S, t->fd, 1); + CHECK(mfd >= 0 && mfd != t->fd); + float a[8], b[8]; + CHECK(pread(t->fd, a, t->nbytes, t->off) == t->nbytes); + CHECK(pread(mfd, b, t->nbytes, t->off) == t->nbytes); + CHECK(memcmp(a, b, sizeof(a)) == 0); + st_prefetch_rep(&S, "t0", 1); /* smoke: WILLNEED on the mirror fd */ + st_prefetch_rep(&S, "t0", 0); + + /* divergent header, same size: rejected */ + CHECK(st_mirror_init(&S, DIR_C) == 0); + CHECK(st_fd_rep(&S, t->fd, 1) == -1); + + /* different size: rejected */ + CHECK(st_mirror_init(&S, DIR_D) == 0); + + /* missing file: rejected (partial mirror with zero shards) */ + CHECK(st_mirror_init(&S, DIR_E) == 0); + + /* unknown fd never maps to a replica */ + CHECK(st_fd_rep(&S, 987654, 1) == -1); + + cleanup(); + puts("safetensors mirror tests: ok"); + return 0; +} diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index c3694c6..ce44fce 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -78,6 +78,15 @@ Format: `VAR` — default — effect. --- +## Dual-SSD streaming + +| Variable | Default | Effect | +|---|---|---| +| `COLI_MODEL_MIRROR` | unset | Path to a second, byte-identical (read-only) copy of the model on another drive; expert reads are split across both. Partial mirrors work (only the shards present are used). | +| `COLI_DISK_WEIGHTS` | unset (startup bandwidth probe) | Split ratio `,` (e.g. `1,1` for 50/50, `9,3` for a fast+slow pair). Unset = probe both drives with the engine's own access pattern at startup. | + +Per-drive byte counts are reported in a `MIRROR:` stats line. Combine with `DIRECT=1` so the two copies never compete for page cache. + ## CUDA (NVIDIA) | Variable | Default | Effect |