Dual-SSD streaming: COLI_MODEL_MIRROR reads experts from two model copies
A second (read-only) copy of the model on another drive is registered as a per-shard read replica; expert loads are split between the drives by a deterministic (layer,eid) hash weighted by COLI_DISK_WEIGHTS=<primary>,<mirror> or, by default, by a startup bandwidth probe using the engine's own access pattern (parallel ~19 MB reads, O_DIRECT). Cold decode is disk-bound (~11 GB/token), so two NVMe queues add up. - st.h: mirror accepted per file only if size + safetensors header are byte-identical to the primary (identical data_offsets by construction, so every pread is valid on either copy); partial mirrors work (smaller second SSD holding only some shards); the mirror is never written — .coli_usage / .coli_kv stay on the primary. - glm.c: routing covers the coalesced slab pread, O_DIRECT, mmap/Metal and scale reads, plus the OMP-parallel pin/autopin warmup (streams from both drives). Deterministic routing keeps readahead/PILOT WILLNEED on the same drive as the demand read and avoids caching an expert twice. A mirror read error falls back to the primary (one warning, no crash). Per-drive bytes are reported in a MIRROR: stats line. - tests: test_st_mirror covers validation, read equality on both replicas, and the rejection paths (divergent header, size mismatch, missing file). Measured on 2x NVMe (GLM-5.2 int4, greedy, DRAFT=0, DIRECT=1, cold-ish cache): 0.42 -> 0.57 tok/s (+36%), expert-disk service 15.2s -> 10.3s, byte-identical output; probe chose a 48/52 split.
This commit is contained in:
@@ -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 <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
|
||||
#include "../st.h"
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <direct.h>
|
||||
#define MKDIR(p) _mkdir(p)
|
||||
#else
|
||||
#include <sys/stat.h>
|
||||
#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;
|
||||
}
|
||||
Reference in New Issue
Block a user