diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1dd58088b..631adf246 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,50 @@ 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 llama_context_restore_device(ctx_tgt); if (ctx_dft != nullptr) { llama_context_restore_device(ctx_dft); @@ -1089,6 +1119,9 @@ 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. Clearing it earlier makes our warden honour rings + // the ringer has already satisfied, so we hand the token straight back without serving. A + // ring genuinely dropped here is harmless: vram_take_token() re-rings about once a second. vram_cold = false; } @@ -1098,6 +1131,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