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:
@@ -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();
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user