diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 53521fc667..fdc8f2b805 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3580,8 +3580,8 @@ private: has_output |= batch.tokens[i].output; } - // decode on the worker thread, so we can still handle metrics tasks while waiting - // note: the sync is done here too, so that the wait also happens off the main thread + // yield to the queue, so we can still handle metrics tasks while decoding + // note: the sync is done here too, so that the wait is also covered by the yield int ret = 0; queue_tasks.yield_to_queue([&]() { ret = llama_decode(ctx_tgt, batch_view); diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index d22527ffec..2bcc9bd8f2 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -150,31 +150,46 @@ bool server_queue::process_new_tasks(bool is_yielding) { void server_queue::worker_loop() { while (true) { - std::function work; { std::unique_lock lock(mutex_tasks); + // wait on busy instead of yielding - busy stays set even when the yield already ended worker.cv.wait(lock, [&]{ - return worker.stop || worker.work != nullptr; + return worker.stop || worker.busy; }); if (worker.stop) { return; } - work = std::move(worker.work); - worker.work = nullptr; } - // note: do not hold any lock here, work() may post new tasks - std::exception_ptr exception; - try { - work(); - } catch (...) { - exception = std::current_exception(); + // process tasks while the yield is active + while (true) { + bool terminated = false; + try { + // note: do not hold any lock here, the callback may post new tasks + terminated = process_new_tasks(true); + } catch (...) { + std::unique_lock lock(mutex_tasks); + worker.exception = std::current_exception(); + break; + } + + std::unique_lock lock(mutex_tasks); + if (terminated || worker.stop || !worker.yielding) { + break; + } + if (!queue_tasks.empty()) { + continue; // a new task arrived in the meantime + } + condition_tasks.wait(lock, [&]{ + return worker.stop || !running || !worker.yielding || !queue_tasks.empty(); + }); } - // signal completion to yield_to_queue() - std::unique_lock lock(mutex_tasks); - worker.exception = std::move(exception); - worker.busy = false; + // signal to yield_to_queue() that no more tasks will be processed + { + std::unique_lock lock(mutex_tasks); + worker.busy = false; + } condition_tasks.notify_all(); } } @@ -188,6 +203,7 @@ void server_queue::worker_stop() { worker.stop = true; } worker.cv.notify_one(); + condition_tasks.notify_all(); worker.thread.join(); } @@ -199,29 +215,29 @@ void server_queue::yield_to_queue(std::function && work) { { std::unique_lock lock(mutex_tasks); GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested"); - worker.busy = true; - worker.work = std::move(work); + worker.busy = true; + worker.yielding = true; } worker.cv.notify_one(); - while (true) { - // note: on terminate this is a no-op, but we still wait for the work to finish - process_new_tasks(true); - - std::unique_lock lock(mutex_tasks); - // declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new - condition_tasks.wait(lock, [&]{ - return !worker.busy || (running && !queue_tasks.empty()); - }); - if (!worker.busy) { - break; - } + // run the work on the current thread, so that all ggml compute stays on the same thread + std::exception_ptr exception; + try { + work(); + } catch (...) { + exception = std::current_exception(); } - std::exception_ptr exception; { std::unique_lock lock(mutex_tasks); + // the yield is over, wait for the worker to finish its current task + worker.yielding = false; + condition_tasks.notify_all(); + condition_tasks.wait(lock, [&]{ + return !worker.busy; + }); + // put the declined tasks back, keeping their order while (!queue_tasks_unhandled.empty()) { queue_tasks.push_front(std::move(queue_tasks_unhandled.back())); @@ -231,8 +247,12 @@ void server_queue::yield_to_queue(std::function && work) { // make sure to avoid idle timeout here time_last_task = ggml_time_ms(); - // the worker is idle now, take the exception it may have left behind - std::swap(exception, worker.exception); + // an exception from work() takes precedence over the one from the worker + if (!exception) { + std::swap(exception, worker.exception); + } else { + worker.exception = nullptr; + } } QUE_DBG("%s", "done yielding to queue\n"); @@ -249,7 +269,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { // spawn the worker thread used by yield_to_queue() GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running"); - worker.stop = false; + worker.stop = false; + worker.busy = false; + worker.yielding = false; worker.thread = std::thread([this]() { worker_loop(); }); constexpr auto max_wait_time = std::chrono::seconds(1); diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 69a3e554f5..52d30095c1 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -33,11 +33,11 @@ private: // used by yield_to_queue, all fields are guarded by mutex_tasks struct worker_t { std::thread thread; - std::condition_variable cv; // the worker sleeps on this until there is work - std::function work; // pending work, picked up by the thread - std::exception_ptr exception; // exception thrown by work(), if any - bool stop = false; - bool busy = false; + std::condition_variable cv; // the worker sleeps on this until a yield starts + std::exception_ptr exception; // exception thrown while processing tasks, if any + bool stop = false; + bool busy = false; // set by yield_to_queue(), cleared by the worker once it is done processing tasks + bool yielding = false; // work() is still running on the start_loop() thread }; worker_t worker; @@ -93,7 +93,7 @@ public: */ void start_loop(int64_t idle_sleep_ms = -1); - // run work() on a separate thread, while the current thread calls process_new_tasks + // while waiting for work() to finish, run process_new_tasks on the worker thread // returns once work() is done (may throw exceptions) // must be called from start_loop() thread (ideally inside callback_update_slots) // use case: return metrics while encode/decode is running @@ -116,6 +116,7 @@ public: // the second argument tells whether the queue is currently yielding (see yield_to_queue) // only then may the callback return false to decline the task, and it must leave it // untouched, so that it can be put back in the queue later + // note: while yielding, the callback runs on worker thread, not main thread void on_new_task(std::function callback) { callback_new_task = std::move(callback); }