Steve Markgraf
63966ba3f8
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.
110 lines
3.6 KiB
C
110 lines
3.6 KiB
C
/* 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;
|
|
}
|