diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1dd58088b..cc69d73c4 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -46,6 +46,7 @@ #include #include #include +#include #endif constexpr int HTTP_POLLING_SECONDS = 1; @@ -988,12 +989,9 @@ 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 } @@ -1068,18 +1066,53 @@ private: #endif } + // take the shared VRAM token: ring the doorbell, then block on the flock. The doorbell is + // edge-triggered and only heard by a process that is already warm and listening, so a ring sent + // while the holder still loads or restores is lost - retry the lock and re-ring about once a + // second until we get it. 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_ring_doorbell(); + for (int tries = 1; ; tries++) { + if (flock(vram_lock_fd, LOCK_EX | LOCK_NB) == 0) { + return true; + } + if (errno == EINTR) { + continue; // our signal handlers run without SA_RESTART + } + if (errno != EWOULDBLOCK && errno != EAGAIN) { + SRV_WRN("VRAM arbiter: flock failed (errno %d)\n", errno); + return false; + } + poll(nullptr, 0, 200); + if (tries % 5 == 0) { + vram_ring_doorbell(); // the holder may have missed the ring while it was not listening + } + if (tries % 25 == 0) { + SRV_WRN("VRAM arbiter: still waiting for the GPU token after %d s\n", tries / 5); + } + } +#else + return true; +#endif + } + // acquire the VRAM token (blocking) and bring weights back to the device. Called on the loop // thread right before a decode, so it never races compute. void vram_ensure_warm() { 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 + // cold means "does not hold the token", so clear it before the slow restore - while it is + // set our own warden drops foreign rings, and a dropped ring is never re-sent + vram_cold = false; llama_context_restore_device(ctx_tgt); if (ctx_dft != nullptr) { llama_context_restore_device(ctx_dft); @@ -1089,7 +1122,6 @@ 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. - vram_cold = false; } void vram_share_shutdown() { @@ -1098,6 +1130,10 @@ private: if (vram_warden.joinable()) { vram_warden.join(); } + // drop our own ring only - a foreign file can still be a live pending request + if (!vram_doorbell_dir.empty() && !vram_pid_str.empty()) { + unlink((vram_doorbell_dir + "/" + vram_pid_str).c_str()); + } 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; } #endif @@ -1118,12 +1154,12 @@ private: if (mctx != nullptr) { mtmd_release_device(mctx); } + vram_cold = true; // drop the token flag first: a waiter can lock the moment we unlock #if !defined(_WIN32) if (vram_flock) { flock(vram_lock_fd, LOCK_UN); } #endif - vram_cold = true; } common_speculative_init_result_ptr spec_init;