url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #331 from nbeerbower/st-pread-full

Browse Source 
st: chunked pread with EINTR retry and honest short-read errors
This commit is contained in:
Vincenzo
2026-07-17 20:42:42 +02:00
committed by GitHub
parent 19f267794f ed8dab4da4
commit 1a956b0119
3 changed files with 131 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/Makefile
+4 -1
View File
@@ -166,7 +166,7 @@ else
PYTHON ?= python3
endif
CUDA_OBJ =
TEST_BINS = tests/test_json$(EXE) tests/test_st$(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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
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_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
ifneq (,$(LINUX))
TEST_BINS += tests/test_uring$(EXE)
endif
@@ -293,6 +293,9 @@ iobench$(EXE): iobench.c compat.h
tests/test_json$(EXE): tests/test_json.c json.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
tests/test_st_pread$(EXE): tests/test_st_pread.c st.h json.h compat.h
$(CC) $(CFLAGS) -DST_PREAD_CHUNK=7 $< -o $@ $(LDFLAGS)
tests/test_st$(EXE): tests/test_st.c st.h json.h compat.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
c/st.h
+38 -5
View File
@@ -12,6 +12,7 @@
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <dirent.h>
#include <sys/stat.h>
@@ -104,6 +105,38 @@ static int st_direct_fd(shards *S, int fd) {
}
/* indicizza tutti i model-*.safetensors in snap_dir */
/* pread completo: chunk-loop (una singola pread si ferma a ~2^31 byte su Linux
* — i tensori bf16 grandi la superano), riprova su EINTR e riporta un errore
* ONESTO: perror stampava "Success" su una short-read (errno resta 0), lo
* stesso sintomo corretto in glm.c per #236. ST_PREAD_CHUNK e' sovrascrivibile
* per i test. EN: full pread — chunk loop (one pread caps at ~2^31 bytes and
* big bf16 tensors exceed it), EINTR retry, honest short-read errors.
* Exits on failure, like every st.h reader. */
#ifndef ST_PREAD_CHUNK
#define ST_PREAD_CHUNK (1u << 30)
#endif
static void st_pread_full(int fd, void *buf, int64_t n, int64_t off, const char *tag) {
char *p = (char *)buf;
int64_t got = 0;
while (got < n) {
int64_t want = n - got;
if (want > (int64_t)ST_PREAD_CHUNK) want = ST_PREAD_CHUNK;
ssize_t r = pread(fd, p + got, (size_t)want, off + got);
if (r < 0) {
if (errno == EINTR) continue;
fprintf(stderr, "%s: %s (off %lld, %lld/%lld bytes)\n", tag, strerror(errno),
(long long)off, (long long)got, (long long)n);
exit(1);
}
if (r == 0) {
fprintf(stderr, "%s: short read at EOF (off %lld, %lld/%lld bytes) — truncated file?\n",
tag, (long long)off, (long long)got, (long long)n);
exit(1);
}
got += r;
}
}
static void st_init(shards *S, const char *snap_dir) {
memset(S, 0, sizeof(*S));
S->cap = 4096; S->t = calloc(S->cap, sizeof(st_tensor));
@@ -128,7 +161,7 @@ static void st_init(shards *S, const char *snap_dir) {
if (fstat(fd, &sst) != 0) { perror("fstat shard"); exit(1); }
int64_t fsz = (int64_t)sst.st_size;
uint64_t hlen;
if (pread(fd, &hlen, 8, 0) != 8) { perror("pread hlen"); exit(1); }
st_pread_full(fd, &hlen, 8, 0, "pread hlen");
/* file malevolo/troncato: hlen deve stare nel file dopo gli 8 byte di
* prefisso e sotto il tetto. Senza questo bound hlen+1 puo' andare in
* overflow (malloc(0) e poi hdr[hlen]=0 fuori limiti) o forzare una
@@ -138,7 +171,7 @@ static void st_init(shards *S, const char *snap_dir) {
files[fi], (unsigned long long)hlen, (long long)fsz); exit(1); }
char *hdr = malloc(hlen + 1);
if (!hdr) { perror("malloc safetensors header"); exit(1); }
if (pread(fd, hdr, hlen, 8) != (ssize_t)hlen) { perror("pread hdr"); exit(1); }
st_pread_full(fd, hdr, (int64_t)hlen, 8, "pread hdr");
hdr[hlen] = 0;
int64_t data_start = 8 + (int64_t)hlen;
char *arena = NULL;
@@ -218,7 +251,7 @@ static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) {
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
void *raw = malloc(t->nbytes);
if (!raw) { fprintf(stderr, "malloc %lld bytes for tensor %s failed\n", (long long)t->nbytes, name); exit(1); }
if (pread(t->fd, raw, t->nbytes, t->off) != t->nbytes) { perror("pread data"); exit(1); }
st_pread_full(t->fd, raw, t->nbytes, t->off, "pread data");
if (t->dtype == 2) {
memcpy(out, raw, t->nbytes);
} else if (t->dtype == 0) {
@@ -243,7 +276,7 @@ static int64_t st_nbytes(shards *S, const char *name) {
static void st_read_raw(shards *S, const char *name, void *out, int drop) {
st_tensor *t = st_find(S, name);
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
if (pread(t->fd, out, t->nbytes, t->off) != t->nbytes) { perror("pread raw"); exit(1); }
st_pread_full(t->fd, out, t->nbytes, t->off, "pread raw");
if (drop) posix_fadvise(t->fd, t->off, t->nbytes, POSIX_FADV_DONTNEED);
}
@@ -256,7 +289,7 @@ static void st_read_slice_f32(shards *S, const char *name, int64_t elem_off, int
int esz = (t->dtype == 2) ? 4 : 2;
int64_t boff = t->off + elem_off * esz, nb = n_elems * esz;
void *raw = malloc(nb);
if (pread(t->fd, raw, nb, boff) != nb) { perror("pread slice"); exit(1); }
st_pread_full(t->fd, raw, nb, boff, "pread slice");
if (t->dtype == 2) memcpy(out, raw, nb);
else if (t->dtype == 0) { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = bf16_to_f32(p[i]); }
else { uint16_t *p = raw; for (int64_t i = 0; i < n_elems; i++) out[i] = f16_to_f32(p[i]); }
c/tests/test_st_pread.c
+89
View File
@@ -0,0 +1,89 @@
/* st_pread_full: chunk loop + honest truncation errors.
* Built with -DST_PREAD_CHUNK=7 so a ~100-byte tensor takes many pread calls —
* exercising the loop that production only needs past 2^31 bytes (one pread
* caps there on Linux; big bf16 tensors exceed it). Also forks a child against
* a truncated shard and requires exit(1) with a "short read" message instead
* of the old perror("... : Success"). */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifndef _WIN32
#include <sys/wait.h>
#include <unistd.h>
#endif
#include "../st.h"
#define CHECK(condition) do { \
if (!(condition)) { \
fprintf(stderr, "%s:%d: check failed: %s\n", __FILE__, __LINE__, #condition); \
return 1; \
} \
} while (0)
static void write_snap(const char *dir, int truncate_bytes) {
char path[512];
snprintf(path, sizeof(path), "%s/model.safetensors", dir);
unsigned char data[96];
for (int i = 0; i < 96; i++) data[i] = (unsigned char)(i * 7 + 3);
const char *hdr = "{\"t\":{\"dtype\":\"U8\",\"shape\":[96],\"data_offsets\":[0,96]}}";
uint64_t hlen = strlen(hdr);
FILE *f = fopen(path, "wb");
fwrite(&hlen, 8, 1, f);
fwrite(hdr, 1, hlen, f);
fwrite(data, 1, (size_t)(96 - truncate_bytes), f);
fclose(f);
}
int main(void) {
/* relative to the CWD, per test_stops: MinGW .exe files resolve Windows
* paths and "/tmp" is not one */
char dir[] = "test_st_pread_XXXXXX";
if (!mkdtemp(dir)) { perror("mkdtemp"); return 1; }
/* 1) chunk loop: 96-byte tensor read 7 bytes at a time, content exact */
write_snap(dir, 0);
shards S; st_init(&S, dir);
unsigned char out[96] = {0};
st_read_raw(&S, "t", out, 0);
for (int i = 0; i < 96; i++) CHECK(out[i] == (unsigned char)(i * 7 + 3));
#ifndef _WIN32
/* 2) shard truncated AFTER st_init (init validates static bounds, so the
* pread path only fires when the file shrinks underneath a live handle):
* child must exit(1) with an honest message, not perror's "Success" */
char shard[512]; snprintf(shard, sizeof(shard), "%s/model.safetensors", dir);
struct stat sb; CHECK(stat(shard, &sb) == 0);
CHECK(truncate(shard, sb.st_size - 40) == 0);
int pipefd[2]; CHECK(pipe(pipefd) == 0);
pid_t pid = fork(); CHECK(pid >= 0);
if (pid == 0) {
dup2(pipefd[1], 2); close(pipefd[0]); close(pipefd[1]);
unsigned char buf[96];
st_read_raw(&S, "t", buf, 0); /* inherited handles; must exit(1) inside */
_exit(42); /* reaching here = bug */
}
close(pipefd[1]);
char err[512] = {0};
ssize_t n = read(pipefd[0], err, sizeof(err)-1); (void)n;
close(pipefd[0]);
int status = 0; waitpid(pid, &status, 0);
CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 1);
CHECK(strstr(err, "short read") != NULL);
CHECK(strstr(err, "Success") == NULL);
#else
/* fork/pipe/truncate are POSIX; Windows still runs the chunk-loop check */
printf("test_st_pread: truncation subtest skipped on Windows\n");
#endif
char cmd[600];
#ifdef _WIN32
snprintf(cmd, sizeof(cmd), "rmdir /s /q %s", dir);
#else
snprintf(cmd, sizeof(cmd), "rm -rf %s", dir);
#endif
if (system(cmd)) {}
printf("test_st_pread: chunk loop + honest truncation error: ok\n");
return 0;
}