diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1dd58088b..031fdaf2d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -38,14 +38,14 @@ #include #endif -// POSIX file locking + inotify doorbell for the cross-process VRAM arbiter (see vram_share_* below) +// POSIX file locking for the cross-process VRAM arbiter (see vram_share_* below) #if !defined(_WIN32) #include #include -#include #include #include #include +#include #endif constexpr int HTTP_POLLING_SECONDS = 1; @@ -898,7 +898,7 @@ public: } ~server_context_impl() { - // stop the VRAM warden thread (and release the token) before tearing anything down + // release the VRAM token before tearing anything down vram_share_shutdown(); if (!sleeping) { // destroy() is already called when entering sleeping state @@ -939,18 +939,20 @@ private: bool vram_only = false; // release-mode sleep active (LLAMA_SLEEP_VRAM_ONLY) bool vram_evict_kv = false; // also evict the KV cache to host on release (LLAMA_SLEEP_EVICT_KV) bool vram_flock = false; // cross-process flock coordination available - std::atomic vram_cold{true}; // weights currently released (read by warden thread) + std::atomic vram_cold{true}; // weights currently released int vram_lock_fd = -1; // fd for /token.lock - // inotify "doorbell": a waiter that wants the VRAM token touches /doorbell/, which - // wakes the current holder's warden thread so it releases immediately instead of only on idle. - std::string vram_doorbell_dir; - std::string vram_pid_str; - int vram_inotify_fd = -1; - std::thread vram_warden; - std::atomic vram_warden_run{false}; + // "want" lock: a process that waits for the token holds /want.lock shared, and drops it + // as soon as it owns the token. The holder probes that lock with a non-blocking exclusive + // flock, which fails only while somebody waits. This is level-triggered state owned by the + // kernel, so there is no event to miss and nothing goes stale if a waiter dies. + // The probe needs its own fd: flock() treats two open file descriptions of one file + // independently, so probing on the waiter fd would convert our own lock instead of conflicting. + int vram_want_fd = -1; // waiter side, LOCK_SH while waiting + int vram_want_probe_fd = -1; // probe side, LOCK_EX | LOCK_NB only + bool vram_want_held = false; // only touched by the waiting thread - // open the shared arena (flock token + doorbell dir). Idempotent; sets vram_only/vram_flock. + // open the shared arena (token.lock + want.lock). Idempotent; sets vram_only/vram_flock. void vram_arena_open() { if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) { return; @@ -968,11 +970,14 @@ private: vram_lock_fd = open(lock_path.c_str(), O_RDWR | O_CREAT, 0666); if (vram_lock_fd >= 0) { vram_flock = true; - vram_pid_str = std::to_string(getpid()); - vram_doorbell_dir = arena + "/doorbell"; - 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()); + const std::string want_path = arena + "/want.lock"; + vram_want_fd = open(want_path.c_str(), O_RDWR | O_CREAT, 0666); + vram_want_probe_fd = open(want_path.c_str(), O_RDWR | O_CREAT, 0666); + if (vram_want_fd < 0 || vram_want_probe_fd < 0) { + SRV_WRN("VRAM arbiter: cannot open %s, the holder will only release on idle\n", want_path.c_str()); + } + SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (want %s)\n", + lock_path.c_str(), want_path.c_str()); } else { SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str()); } @@ -988,18 +993,15 @@ private: 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 + if (!vram_take_token()) { + SRV_WRN("%s", "VRAM arbiter: loading weights without the GPU token\n"); } -#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. + // 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 queue + // releases it once another model registers on want.lock. void vram_share_init() { vram_arena_open(); if (!vram_only) { @@ -1007,64 +1009,92 @@ private: } 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); - if (vram_inotify_fd >= 0) { - inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE); - vram_warden_run = true; - vram_warden = std::thread([this]{ vram_warden_loop(); }); - } + if (vram_flock && vram_want_probe_fd >= 0) { + queue_tasks.on_should_yield([this]{ return vram_should_yield(); }); } #endif } - // warden thread: block on the doorbell; when another process rings (wants the token) and we - // currently hold it, ask the loop to yield (release) at its next idle point. Only sets a flag on - // the queue - it never posts a task and never touches the GPU. The flag is consumed by - // should_sleep() on the start_loop() thread, which is only reached after callback_update_slots() - // has returned, so a release can never happen underneath an in-flight decode - not even via the - // worker thread of yield_to_queue(), which declines everything except read-only tasks. - void vram_warden_loop() { + // true while another process waits for the token and we still hold it. Called by should_sleep() + // on the start_loop() thread, under mutex_tasks, and only when the queue is idle - so a release + // can never happen underneath an in-flight decode, and there is no event to latch: the answer + // comes from live kernel state at the moment the loop asks. + bool vram_should_yield() { #if !defined(_WIN32) - char buf[4096]; - while (vram_warden_run.load()) { - struct pollfd pfd { vram_inotify_fd, POLLIN, 0 }; - int pr = poll(&pfd, 1, 500); // 500ms so we periodically re-check the run flag - if (pr <= 0) { - continue; - } - ssize_t n = read(vram_inotify_fd, buf, sizeof(buf)); - if (n <= 0) { - continue; - } - bool foreign_ring = false; - for (char * p = buf; p < buf + n; ) { - struct inotify_event * ev = (struct inotify_event *) p; - if (ev->len > 0 && vram_pid_str != ev->name) { - foreign_ring = true; // someone else wants the GPU - } - p += sizeof(struct inotify_event) + ev->len; - } - if (foreign_ring && !vram_cold) { - queue_tasks.request_yield(); - } + if (!vram_flock || vram_cold || vram_want_probe_fd < 0) { + return false; // we hold no token to give away } + // never blocks: LOCK_NB only, and we drop the probe lock at once if we win it + if (flock(vram_want_probe_fd, LOCK_EX | LOCK_NB) == 0) { + flock(vram_want_probe_fd, LOCK_UN); + return false; // nobody wants the GPU + } + return errno == EWOULDBLOCK || errno == EAGAIN; +#else + return false; #endif } - // ring the doorbell so the current token holder releases promptly - void vram_ring_doorbell() { + // register as a waiter, so the current token holder sees us and yields. Shared, so several + // waiters coexist. The kernel drops it if we die, so it can never strand the holder. + void vram_want_acquire() { #if !defined(_WIN32) - if (!vram_flock || vram_doorbell_dir.empty()) { + if (vram_want_fd < 0 || vram_want_held) { return; } - const std::string f = vram_doorbell_dir + "/" + vram_pid_str; - int fd = open(f.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666); - if (fd >= 0) { - ssize_t w = write(fd, "1", 1); - (void) w; - close(fd); + while (flock(vram_want_fd, LOCK_SH) != 0) { + if (errno != EINTR) { // our signal handlers run without SA_RESTART + SRV_WRN("VRAM arbiter: cannot take want.lock (errno %d)\n", errno); + return; + } } + vram_want_held = true; +#endif + } + + // stop asking for the token. Must run on every path out of vram_take_token(), or the holder + // yields forever. + void vram_want_release() { +#if !defined(_WIN32) + if (vram_want_fd < 0 || !vram_want_held) { + return; + } + flock(vram_want_fd, LOCK_UN); + vram_want_held = false; +#endif + } + + // take the shared VRAM token: register on want.lock, then wait for token.lock. want.lock is + // level-triggered state, so the holder still learns of us if it is loading or restoring when we + // start to wait. No timeout: the holder can be mid-generation, and decoding without the token + // risks an OOM that kills both processes. + bool vram_take_token() { +#if !defined(_WIN32) + if (!vram_flock) { + return true; // no cross-process lock: there is no token to take + } + vram_want_acquire(); + bool logged = false; + while (flock(vram_lock_fd, LOCK_EX | LOCK_NB) != 0) { + if (errno == EWOULDBLOCK || errno == EAGAIN) { + if (!logged) { + SRV_INF("%s", "VRAM arbiter: waiting for the GPU token\n"); + logged = true; + } + flock(vram_lock_fd, LOCK_EX); // blocking: the kernel grants it the moment the holder unlocks + continue; // re-check with LOCK_NB, which also succeeds on a lock we already own + } + if (errno == EINTR) { + continue; // our signal handlers run without SA_RESTART + } + SRV_WRN("VRAM arbiter: flock failed (errno %d)\n", errno); + vram_want_release(); + return false; + } + vram_want_release(); + return true; +#else + return true; #endif } @@ -1074,12 +1104,9 @@ private: if (!vram_only || !vram_cold) { return; } -#if !defined(_WIN32) - if (vram_flock) { - vram_ring_doorbell(); // nudge the current holder to release - flock(vram_lock_fd, LOCK_EX); // blocks until the current holder goes cold + if (!vram_take_token()) { + SRV_WRN("%s", "VRAM arbiter: restoring weights without the GPU token\n"); } -#endif llama_context_restore_device(ctx_tgt); if (ctx_dft != nullptr) { llama_context_restore_device(ctx_dft); @@ -1089,17 +1116,18 @@ private: // on-demand mode, released again right after - so a warm text-only model holds no encoder // VRAM (that space is free for KV / expert cache). A cold->warm wake for a *text* request // therefore leaves the encoder in RAM; an image request restores it at encode time. + // NOTE: keep this AFTER the restore, so the loop cannot give away the token while we still + // bring the weights up. We dropped want.lock when we took the token, so the next probe + // reports only the waiters that are still there. vram_cold = false; } void vram_share_shutdown() { #if !defined(_WIN32) - vram_warden_run = false; - if (vram_warden.joinable()) { - vram_warden.join(); - } - if (vram_inotify_fd >= 0) { close(vram_inotify_fd); vram_inotify_fd = -1; } - if (vram_lock_fd >= 0) { flock(vram_lock_fd, LOCK_UN); close(vram_lock_fd); vram_lock_fd = -1; } + vram_want_release(); + if (vram_want_probe_fd >= 0) { close(vram_want_probe_fd); vram_want_probe_fd = -1; } + if (vram_want_fd >= 0) { close(vram_want_fd); vram_want_fd = -1; } + if (vram_lock_fd >= 0) { flock(vram_lock_fd, LOCK_UN); close(vram_lock_fd); vram_lock_fd = -1; } #endif } @@ -1688,9 +1716,9 @@ private: handle_sleeping_state(sleeping); }); - // VRAM arbiter: start the doorbell warden. The load was coordinated (vram_acquire_for_load - // grabbed the token before uploading), so we stay warm holding the token and serve the first - // request without a re-warm; the warden releases us when another model rings the doorbell. + // VRAM arbiter: arm the yield probe. The load was coordinated (vram_acquire_for_load + // grabbed the token before uploading), so we stay warm holding the token and serve the + // first request without a re-warm; the loop releases it when another model waits for it. vram_share_init(); metrics.init(); @@ -4514,9 +4542,9 @@ struct server_res_generator : server_res_spipe { server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false) : rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) { // fast path in case sleeping is disabled. Note: the VRAM arbiter (LLAMA_SLEEP_VRAM_ONLY) - // can put the server to sleep via the cross-process doorbell even when idle-sleep is - // disabled (sleep_idle_seconds < 0), so in that case requests must still wake it - otherwise - // a doorbell-slept server would hang, never returning from a request. + // can put the server to sleep to hand over the GPU token even when idle-sleep is disabled + // (sleep_idle_seconds < 0), so in that case requests must still wake it - otherwise an + // arbiter-slept server would hang, never returning from a request. static const bool vram_arbiter = getenv("LLAMA_SLEEP_VRAM_ONLY") != nullptr; bypass_sleep |= (sleep_idle_seconds < 0 && !vram_arbiter); if (!bypass_sleep) { diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index f173043f0..b7ce86433 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -114,27 +114,19 @@ void server_queue::pop_deferred_task(int id_slot) { void server_queue::wait_until_no_sleep() { std::unique_lock lock(mutex_tasks); - if (!sleeping) { - return; - } else { + // re-ask on every wake: the loop clears req_stop_sleeping on the way into sleep, so a loop that + // goes back to sleep before we run would strand us here forever + while (sleeping) { if (!req_stop_sleeping) { QUE_DBG("%s", "requesting to stop sleeping\n"); req_stop_sleeping = true; - condition_tasks.notify_one(); // only main thread is waiting on this + condition_tasks.notify_all(); // other threads may wait on this too } QUE_DBG("%s", "waiting until no sleep\n"); - condition_tasks.wait(lock, [&]{ - return !sleeping; - }); + condition_tasks.wait(lock); } } -void server_queue::request_yield() { - std::unique_lock lock(mutex_tasks); - yield_requested = true; - condition_tasks.notify_all(); -} - void server_queue::terminate() { std::unique_lock lock(mutex_tasks); running = false; @@ -292,11 +284,17 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { worker.yielding = false; worker.thread = std::thread([this]() { worker_loop(); }); - constexpr auto max_wait_time = std::chrono::seconds(1); + // the arbiter predicate is read from live state, so this timeout bounds how fast we notice a + // process waiting for the VRAM token + const auto max_wait_time = should_yield_cb ? std::chrono::milliseconds(10) : std::chrono::milliseconds(1000); + // after a wake, keep the VRAM for at least this long: the request that woke us is not in the + // queue yet, so yielding at once would only send it through another wake + constexpr int64_t yield_grace_ms = 100; + int64_t time_last_wake = 0; auto should_sleep = [&]() -> bool { // caller must hold mutex_tasks - if (yield_requested) { - return true; // another process rang the VRAM doorbell - release now + if (should_yield_cb && ggml_time_ms() - time_last_wake >= yield_grace_ms && should_yield_cb()) { + return true; // another process waits for the VRAM token - release now } if (idle_sleep_ms < 0) { return false; @@ -336,15 +334,16 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { if (should_sleep()) { QUE_INF("%s", "entering sleeping state\n"); sleeping = true; - yield_requested = false; // consumed // Call order cb0 -> cb1 -> cb{N} for (auto & cb : callback_sleeping_state) { cb(true); } req_stop_sleeping = false; - // wait until we are requested to exit sleeping state + // wait until we are requested to exit sleeping state, or a task arrives: post() only + // notifies, so a task queued right after wait_until_no_sleep() saw us awake must be + // able to wake us by itself, else it waits here for an unrelated request condition_tasks.wait(lock, [&]{ - return (!running || req_stop_sleeping); + return (!running || req_stop_sleeping || !queue_tasks.empty()); }); if (!running) { // may changed during sleep break; // terminate @@ -357,12 +356,13 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { } sleeping = false; time_last_task = ggml_time_ms(); + time_last_wake = time_last_task; condition_tasks.notify_all(); // notify wait_until_no_sleep() break; // process new tasks } else { - // wait for new tasks, a VRAM yield request, or timeout for checking sleeping condition + // wait for new tasks, or timeout for checking sleeping condition bool res = condition_tasks.wait_for(lock, max_wait_time, [&]{ - return (!queue_tasks.empty() || !running || yield_requested); + return (!queue_tasks.empty() || !running); }); if (res && !queue_tasks.empty()) { break; // new task arrived or terminate @@ -370,7 +370,7 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { if (!running) { break; } - // otherwise (timeout or yield request), loop again to re-check should_sleep + // otherwise (timeout), loop again to re-check should_sleep } } } diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 4bc390c93..7939cb299 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -18,7 +18,6 @@ private: bool running = false; bool sleeping = false; bool req_stop_sleeping = false; - bool yield_requested = false; // set by request_yield() when another process wants the VRAM token int64_t time_last_task = 0; // queues @@ -46,6 +45,7 @@ private: std::function callback_new_task; std::function callback_update_slots; std::vector> callback_sleeping_state; + std::function should_yield_cb; public: ~server_queue() { worker_stop(); } @@ -70,11 +70,6 @@ public: // returns immediately if not sleeping void wait_until_no_sleep(); - // request that the loop go to sleep (release VRAM) as soon as it is idle - called from the - // VRAM-arbiter warden thread when another process rings the doorbell for the GPU token. - // Thread-safe; wakes the loop so it releases promptly instead of at the next idle poll. - void request_yield(); - bool is_sleeping() { std::unique_lock lock(mutex_tasks); return sleeping; @@ -133,6 +128,14 @@ public: callback_update_slots = std::move(callback); } + // Register a predicate asking the loop to sleep (release VRAM) as soon as it is idle - used by + // the VRAM arbiter to check whether another process waits for the GPU token. + // Called on the start_loop() thread while holding mutex_tasks, so it must not block or post + // tasks. While it is set, the idle wait polls faster, so the loop reacts without a notify. + void on_should_yield(std::function callback) { + should_yield_cb = std::move(callback); + } + // Register callback for sleeping state change; multiple callbacks are allowed // for example: register order cb0, cb1, cb2 // entering sleep: queue.sleeping = true --> cb0(true) --> cb1(true) --> cb2(true)