server: coordinate model load with the VRAM arbiter to avoid load-time OOM

Before uploading a model's weights, acquire the shared VRAM token (ring the
doorbell so any resident model releases first). Previously load uploaded weights
to VRAM before the arbiter was active, so loading a large model while another
(e.g. the warm 4B task model) held VRAM could exceed the budget and OOM.

- add vram_arena_open() (idempotent flock/doorbell setup) and
  vram_acquire_for_load(), called from load_model() before
  common_init_from_params().
- a coordinated load now stays warm holding the token and serves its first
  request without a re-warm (drop the init-time go_cold cold-start).

Assisted-by: Claude
This commit is contained in:
2026-07-25 14:36:16 +02:00
parent bf68c5446e
commit f2d79d89f5
+50 -10
View File
@@ -927,14 +927,17 @@ private:
std::thread vram_warden; std::thread vram_warden;
std::atomic<bool> vram_warden_run{false}; std::atomic<bool> vram_warden_run{false};
void vram_share_init() { // open the shared arena (flock token + doorbell dir). Idempotent; sets vram_only/vram_flock.
void vram_arena_open() {
if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) { if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) {
return; return;
} }
vram_only = true; vram_only = true;
vram_evict_kv = getenv("LLAMA_SLEEP_EVICT_KV") != nullptr; vram_evict_kv = getenv("LLAMA_SLEEP_EVICT_KV") != nullptr;
vram_cold = false; // weights are resident right after load
#if !defined(_WIN32) #if !defined(_WIN32)
if (vram_lock_fd >= 0) {
return; // already open
}
const char * arena_env = getenv("LLAMA_VRAM_ARENA"); const char * arena_env = getenv("LLAMA_VRAM_ARENA");
const std::string arena = arena_env ? arena_env : "/dev/shm/llama-vram"; const std::string arena = arena_env ? arena_env : "/dev/shm/llama-vram";
mkdir(arena.c_str(), 0777); mkdir(arena.c_str(), 0777);
@@ -945,16 +948,49 @@ private:
vram_pid_str = std::to_string(getpid()); vram_pid_str = std::to_string(getpid());
vram_doorbell_dir = arena + "/doorbell"; vram_doorbell_dir = arena + "/doorbell";
mkdir(vram_doorbell_dir.c_str(), 0777); mkdir(vram_doorbell_dir.c_str(), 0777);
SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (doorbell %s)\n",
lock_path.c_str(), vram_doorbell_dir.c_str());
} else {
SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str());
}
#endif
}
// Acquire the VRAM token BEFORE uploading this model's weights, so any model currently resident
// on the GPU releases first and the load uploads into free VRAM instead of racing it (which
// could OOM, e.g. the 4B task model holding VRAM while a large model loads). Blocks until free.
// Called from load_model() right before common_init_from_params().
void vram_acquire_for_load() {
vram_arena_open();
if (!vram_only) {
return;
}
#if !defined(_WIN32)
if (vram_flock) {
vram_ring_doorbell(); // nudge whoever is resident to release
flock(vram_lock_fd, LOCK_EX); // block until the GPU is free
}
#endif
vram_cold = false; // we hold the token; weights will be resident after the upload
}
// Start the doorbell warden. Called from init() after the (coordinated) load. The model stays
// warm (holding the token acquired in vram_acquire_for_load) and serves its first request
// without a re-warm; the warden releases it when another model rings the doorbell.
void vram_share_init() {
vram_arena_open();
if (!vram_only) {
return;
}
vram_cold = false; // resident and holding the token after the coordinated load
#if !defined(_WIN32)
if (vram_flock && vram_inotify_fd < 0) {
vram_inotify_fd = inotify_init1(IN_NONBLOCK); vram_inotify_fd = inotify_init1(IN_NONBLOCK);
if (vram_inotify_fd >= 0) { if (vram_inotify_fd >= 0) {
inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE); inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE);
vram_warden_run = true; vram_warden_run = true;
vram_warden = std::thread([this]{ vram_warden_loop(); }); vram_warden = std::thread([this]{ vram_warden_loop(); });
} }
SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (doorbell %s)\n",
lock_path.c_str(), vram_doorbell_dir.c_str());
} else {
SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str());
} }
#endif #endif
} }
@@ -1321,6 +1357,11 @@ private:
params_base.load_progress_callback_user_data = &load_progress_text; params_base.load_progress_callback_user_data = &load_progress_text;
} }
// VRAM arbiter: acquire the GPU token before uploading weights, so any resident model (e.g.
// the warm 4B task model) releases first and this load uploads into free VRAM instead of
// racing it and OOM-ing. No-op unless LLAMA_SLEEP_VRAM_ONLY is set.
vram_acquire_for_load();
llama_init = common_init_from_params(params_base); llama_init = common_init_from_params(params_base);
model_tgt = llama_init->model(); model_tgt = llama_init->model();
@@ -1588,11 +1629,10 @@ private:
handle_sleeping_state(sleeping); handle_sleeping_state(sleeping);
}); });
// VRAM arbiter: enable cross-process GPU time-sharing (if LLAMA_SLEEP_VRAM_ONLY is set) and // VRAM arbiter: start the doorbell warden. The load was coordinated (vram_acquire_for_load
// start cold - release the just-loaded device weights and drop the shared token, so we only // grabbed the token before uploading), so we stay warm holding the token and serve the first
// occupy VRAM while actually serving. The first decode re-acquires the token and restores. // request without a re-warm; the warden releases us when another model rings the doorbell.
vram_share_init(); vram_share_init();
vram_go_cold();
metrics.init(); metrics.init();