Felipe97/llama-cpp-compiled
01.1k
1#include "server-task.h"2#include "server-queue.h"3 4#include "log.h"5 6#include <algorithm>7#include <chrono>8#include <thread>9 10#define QUE_INF(fmt, ...) LOG_INF("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)11#define QUE_WRN(fmt, ...) LOG_WRN("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)12#define QUE_ERR(fmt, ...) LOG_ERR("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)13#define QUE_DBG(fmt, ...) LOG_DBG("que %12.*s: " fmt, 12, __func__, __VA_ARGS__)14 15#define RES_INF(fmt, ...) LOG_INF("res %12.*s: " fmt, 12, __func__, __VA_ARGS__)16#define RES_WRN(fmt, ...) LOG_WRN("res %12.*s: " fmt, 12, __func__, __VA_ARGS__)17#define RES_ERR(fmt, ...) LOG_ERR("res %12.*s: " fmt, 12, __func__, __VA_ARGS__)18#define RES_DBG(fmt, ...) LOG_DBG("res %12.*s: " fmt, 12, __func__, __VA_ARGS__)19 20//21// server_queue22//23 24static bool task_resets_idle_timer(server_task_type type) {25 return type != SERVER_TASK_TYPE_METRICS;26}27 28int server_queue::post(server_task && task, bool front) {29 std::unique_lock<std::mutex> lock(mutex_tasks);30 GGML_ASSERT(task.id != -1);31 // if this is cancel task make sure to clean up pending tasks32 if (task.type == SERVER_TASK_TYPE_CANCEL) {33 cleanup_pending_task(task.id_target);34 }35 const int task_id = task.id;36 const bool reset_timer = task_resets_idle_timer(task.type);37 QUE_DBG("new task, id = %d, front = %d\n", task_id, front);38 if (front) {39 queue_tasks.push_front(std::move(task));40 } else {41 queue_tasks.push_back(std::move(task));42 }43 if (reset_timer) {44 time_last_task = ggml_time_ms();45 }46 condition_tasks.notify_one();47 return task_id;48}49 50int server_queue::post(std::vector<server_task> && tasks, bool front) {51 std::unique_lock<std::mutex> lock(mutex_tasks);52 bool reset_timer = false;53 for (auto & task : tasks) {54 if (task.id == -1) {55 task.id = id++;56 }57 // if this is cancel task make sure to clean up pending tasks58 if (task.type == SERVER_TASK_TYPE_CANCEL) {59 cleanup_pending_task(task.id_target);60 }61 reset_timer |= task_resets_idle_timer(task.type);62 QUE_DBG("new task, id = %d/%d, front = %d\n", task.id, (int) tasks.size(), front);63 if (front) {64 queue_tasks.push_front(std::move(task));65 } else {66 queue_tasks.push_back(std::move(task));67 }68 }69 if (reset_timer) {70 time_last_task = ggml_time_ms();71 }72 condition_tasks.notify_one();73 return 0;74}75 76void server_queue::defer(server_task && task) {77 std::unique_lock<std::mutex> lock(mutex_tasks);78 QUE_DBG("defer task, id = %d\n", task.id);79 queue_tasks_deferred.push_back(std::move(task));80 time_last_task = ggml_time_ms();81 condition_tasks.notify_one();82}83 84int server_queue::get_new_id() {85 std::unique_lock<std::mutex> lock(mutex_tasks);86 int new_id = id++;87 return new_id;88}89 90void server_queue::pop_deferred_task(int id_slot) {91 std::unique_lock<std::mutex> lock(mutex_tasks);92 if (!queue_tasks_deferred.empty()) {93 // try to find a task that uses the specified slot94 bool found = false;95 for (auto it = queue_tasks_deferred.begin(); it != queue_tasks_deferred.end(); ++it) {96 if (it->id_slot == id_slot) {97 QUE_DBG("pop deferred task (use slot %d), id_task = %d\n", id_slot, it->id);98 queue_tasks.emplace_front(std::move(*it));99 queue_tasks_deferred.erase(it);100 found = true;101 break;102 }103 }104 // if not tasks found using the slot, just pop the first deferred task (default behavior)105 if (!found) {106 QUE_DBG("pop deferred task, id_task = %d\n", queue_tasks_deferred.front().id);107 queue_tasks.emplace_front(std::move(queue_tasks_deferred.front()));108 queue_tasks_deferred.pop_front();109 }110 }111 time_last_task = ggml_time_ms();112 condition_tasks.notify_one();113}114 115void server_queue::wait_until_no_sleep() {116 std::unique_lock<std::mutex> lock(mutex_tasks);117 if (!sleeping) {118 return;119 } else {120 if (!req_stop_sleeping) {121 QUE_DBG("%s", "requesting to stop sleeping\n");122 req_stop_sleeping = true;123 condition_tasks.notify_one(); // only main thread is waiting on this124 }125 QUE_DBG("%s", "waiting until no sleep\n");126 condition_tasks.wait(lock, [&]{127 return !sleeping;128 });129 }130}131 132void server_queue::terminate() {133 std::unique_lock<std::mutex> lock(mutex_tasks);134 running = false;135 condition_tasks.notify_all();136}137 138bool server_queue::process_new_tasks(bool is_yielding) {139 while (true) {140 std::unique_lock<std::mutex> lock(mutex_tasks);141 if (!running) {142 QUE_DBG("%s", "terminate\n");143 return true;144 }145 if (queue_tasks.empty()) {146 return false;147 }148 server_task task = std::move(queue_tasks.front());149 queue_tasks.pop_front();150 lock.unlock();151 152 QUE_DBG("processing task, id = %d\n", task.id);153 if (!callback_new_task(std::move(task), is_yielding)) {154 // set it aside, do not put it back in the queue, else we offer it again in a loop155 GGML_ASSERT(is_yielding && "a task can only be declined while yielding");156 QUE_DBG("task declined, id = %d\n", task.id);157 lock.lock();158 queue_tasks_unhandled.push_back(std::move(task));159 }160 }161}162 163void server_queue::worker_loop() {164 while (true) {165 {166 std::unique_lock<std::mutex> lock(mutex_tasks);167 // wait on busy instead of yielding - busy stays set even when the yield already ended168 worker.cv.wait(lock, [&]{169 return worker.stop || worker.busy;170 });171 if (worker.stop) {172 return;173 }174 }175 176 // process tasks while the yield is active177 while (true) {178 bool terminated = false;179 try {180 // note: do not hold any lock here, the callback may post new tasks181 terminated = process_new_tasks(true);182 } catch (...) {183 std::unique_lock<std::mutex> lock(mutex_tasks);184 worker.exception = std::current_exception();185 break;186 }187 188 std::unique_lock<std::mutex> lock(mutex_tasks);189 if (terminated || worker.stop || !worker.yielding) {190 break;191 }192 if (!queue_tasks.empty()) {193 continue; // a new task arrived in the meantime194 }195 condition_tasks.wait(lock, [&]{196 return worker.stop || !running || !worker.yielding || !queue_tasks.empty();197 });198 }199 200 // signal to yield_to_queue() that no more tasks will be processed201 {202 std::unique_lock<std::mutex> lock(mutex_tasks);203 worker.busy = false;204 }205 condition_tasks.notify_all();206 }207}208 209void server_queue::worker_stop() {210 if (!worker.thread.joinable()) {211 return;212 }213 {214 std::unique_lock<std::mutex> lock(mutex_tasks);215 worker.stop = true;216 }217 worker.cv.notify_one();218 condition_tasks.notify_all();219 worker.thread.join();220}221 222void server_queue::yield_to_queue(std::function<void()> && work) {223 GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running");224 225 QUE_DBG("%s", "yielding to queue\n");226 227 {228 std::unique_lock<std::mutex> lock(mutex_tasks);229 GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");230 worker.busy = true;231 worker.yielding = true;232 }233 worker.cv.notify_one();234 235 // run the work on the current thread, so that all ggml compute stays on the same thread236 std::exception_ptr exception;237 try {238 work();239 } catch (...) {240 exception = std::current_exception();241 }242 243 {244 std::unique_lock<std::mutex> lock(mutex_tasks);245 246 // the yield is over, wait for the worker to finish its current task247 worker.yielding = false;248 condition_tasks.notify_all();249 condition_tasks.wait(lock, [&]{250 return !worker.busy;251 });252 253 // put the declined tasks back, keeping their order254 while (!queue_tasks_unhandled.empty()) {255 queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));256 queue_tasks_unhandled.pop_back();257 }258 259 // make sure to avoid idle timeout here260 time_last_task = ggml_time_ms();261 262 // an exception from work() takes precedence over the one from the worker263 if (!exception) {264 std::swap(exception, worker.exception);265 } else {266 worker.exception = nullptr;267 }268 }269 270 QUE_DBG("%s", "done yielding to queue\n");271 272 // note: rethrow only after the declined tasks are back in the queue, so they are not lost273 if (exception) {274 std::rethrow_exception(exception);275 }276}277 278void server_queue::start_loop(int64_t idle_sleep_ms) {279 running = true;280 time_last_task = ggml_time_ms();281 282 // spawn the worker thread used by yield_to_queue()283 GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");284 worker.stop = false;285 worker.busy = false;286 worker.yielding = false;287 worker.thread = std::thread([this]() { worker_loop(); });288 289 constexpr auto max_wait_time = std::chrono::seconds(1);290 auto should_sleep = [&]() -> bool {291 // caller must hold mutex_tasks292 if (idle_sleep_ms < 0) {293 return false;294 }295 int64_t now = ggml_time_ms();296 return (now - time_last_task) >= idle_sleep_ms;297 };298 299 while (true) {300 QUE_DBG("%s", "processing new tasks\n");301 if (process_new_tasks(false)) {302 break; // terminate303 }304 305 // all tasks in the current loop is processed, slots data is now ready306 QUE_DBG("%s", "update slots\n");307 308 // this will run the main inference process for all slots309 const int64_t t_update_slots = ggml_time_ms();310 callback_update_slots();311 {312 // update_slots() may take a while to finish, we need to make sure it's not counted as idle313 // shift instead of reset, so that non-task_resets_idle_timer tasks do not delay the sleep314 std::unique_lock<std::mutex> lock(mutex_tasks);315 const int64_t now = ggml_time_ms();316 time_last_task = std::min(now, time_last_task + (now - t_update_slots));317 }318 319 QUE_DBG("%s", "waiting for new tasks\n");320 while (true) {321 std::unique_lock<std::mutex> lock(mutex_tasks);322 if (!running || !queue_tasks.empty()) {323 break; // go back to process new tasks or terminate324 }325 326 // no tasks, check for sleeping state327 if (should_sleep()) {328 QUE_INF("%s", "entering sleeping state\n");329 sleeping = true;330 // Call order cb0 -> cb1 -> cb{N}331 for (auto & cb : callback_sleeping_state) {332 cb(true);333 }334 req_stop_sleeping = false;335 // wait until we are requested to exit sleeping state336 condition_tasks.wait(lock, [&]{337 return (!running || req_stop_sleeping);338 });339 if (!running) { // may changed during sleep340 break; // terminate341 }342 QUE_INF("%s", "exiting sleeping state\n");343 req_stop_sleeping = false;344 // Call order cb{N} -> cb1 -> cb0345 for (size_t i = callback_sleeping_state.size(); i > 0; i--) {346 callback_sleeping_state[i - 1](false);347 }348 sleeping = false;349 time_last_task = ggml_time_ms();350 condition_tasks.notify_all(); // notify wait_until_no_sleep()351 break; // process new tasks352 } else {353 // wait for new tasks or timeout for checking sleeping condition354 bool res = condition_tasks.wait_for(lock, max_wait_time, [&]{355 return (!queue_tasks.empty() || !running);356 });357 if (res) {358 break; // new task arrived or terminate359 }360 // otherwise, loop again to check sleeping condition361 }362 }363 }364 365 worker_stop();366}367 368void server_queue::cleanup_pending_task(int id_target) {369 // no need lock because this is called exclusively by post()370 auto rm_func = [id_target](const server_task & task) {371 return task.id == id_target;372 };373 queue_tasks.erase(374 std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func),375 queue_tasks.end());376 queue_tasks_deferred.erase(377 std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func),378 queue_tasks_deferred.end());379 // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled380 queue_tasks_unhandled.erase(381 std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func),382 queue_tasks_unhandled.end());383}384 385//386// server_response387//388 389void server_response::add_waiting_task_id(int id_task) {390 RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size());391 392 std::unique_lock<std::mutex> lock(mutex_results);393 waiting_task_ids.insert(id_task);394}395 396void server_response::add_waiting_task_ids(const std::unordered_set<int> & id_tasks) {397 std::unique_lock<std::mutex> lock(mutex_results);398 399 for (const auto & id_task : id_tasks) {400 RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size());401 waiting_task_ids.insert(id_task);402 }403}404 405void server_response::remove_waiting_task_id(int id_task) {406 RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size());407 408 std::unique_lock<std::mutex> lock(mutex_results);409 waiting_task_ids.erase(id_task);410 // make sure to clean up all pending results411 queue_results.erase(412 std::remove_if(queue_results.begin(), queue_results.end(), [id_task](const server_task_result_ptr & res) {413 return res->id == id_task;414 }),415 queue_results.end());416}417 418void server_response::remove_waiting_task_ids(const std::unordered_set<int> & id_tasks) {419 std::unique_lock<std::mutex> lock(mutex_results);420 421 for (const auto & id_task : id_tasks) {422 RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size());423 waiting_task_ids.erase(id_task);424 }425}426 427server_task_result_ptr server_response::recv(const std::unordered_set<int> & id_tasks) {428 while (true) {429 std::unique_lock<std::mutex> lock(mutex_results);430 condition_results.wait(lock, [&]{431 if (!running) {432 RES_DBG("%s : queue result stop\n", "recv");433 std::terminate(); // we cannot return here since the caller is HTTP code434 }435 return !queue_results.empty();436 });437 438 for (size_t i = 0; i < queue_results.size(); i++) {439 if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) {440 server_task_result_ptr res = std::move(queue_results[i]);441 queue_results.erase(queue_results.begin() + i);442 return res;443 }444 }445 }446 447 // should never reach here448}449 450server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout) {451 while (true) {452 std::unique_lock<std::mutex> lock(mutex_results);453 454 for (int i = 0; i < (int) queue_results.size(); i++) {455 if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) {456 server_task_result_ptr res = std::move(queue_results[i]);457 queue_results.erase(queue_results.begin() + i);458 return res;459 }460 }461 462 std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout));463 if (!running) {464 RES_DBG("%s : queue result stop\n", __func__);465 std::terminate(); // we cannot return here since the caller is HTTP code466 }467 if (cr_res == std::cv_status::timeout) {468 return nullptr;469 }470 }471 472 // should never reach here473}474 475server_task_result_ptr server_response::recv(int id_task) {476 std::unordered_set<int> id_tasks = {id_task};477 return recv(id_tasks);478}479 480void server_response::send(server_task_result_ptr && result) {481 RES_DBG("sending result for task id = %d\n", result->id);482 483 std::unique_lock<std::mutex> lock(mutex_results);484 for (const auto & id_task : waiting_task_ids) {485 if (result->id == id_task) {486 RES_DBG("task id = %d pushed to result queue\n", result->id);487 488 queue_results.emplace_back(std::move(result));489 condition_results.notify_all();490 return;491 }492 }493}494 495void server_response::broadcast(server_task_result_ptr && result) {496 std::unique_lock<std::mutex> lock(mutex_results);497 for (const auto & id_task : waiting_task_ids) {498 RES_DBG("task id = %d pushed to result queue\n", id_task);499 server_task_result_ptr res_copy(result->clone());500 res_copy->id = id_task; // override id with target task id501 queue_results.emplace_back(std::move(res_copy));502 }503 condition_results.notify_all();504}505 506void server_response::terminate() {507 running = false;508 condition_results.notify_all();509}510 511//512// server_response_reader513//514 515void server_response_reader::post_task(server_task && task, bool front) {516 GGML_ASSERT(id_tasks.empty() && "post_task() can only be called once per reader");517 GGML_ASSERT(!task.is_parent() && "not supported, use post_tasks() instead");518 task.index = 0;519 id_tasks.insert(task.id);520 states.push_back(task.create_state());521 queue_results.add_waiting_task_id(task.id);522 queue_tasks.post(std::move(task), front);523}524 525void server_response_reader::post_tasks(std::vector<server_task> && tasks, bool front) {526 GGML_ASSERT(id_tasks.empty() && "post_tasks() can only be called once per reader");527 id_tasks = server_task::get_list_id(tasks);528 states.reserve(tasks.size());529 size_t index = 0;530 for (auto & task : tasks) {531 task.index = index++;532 states.push_back(task.create_state());533 // for child tasks534 for (auto & child_task : task.child_tasks) {535 child_task.index = index++;536 states.push_back(child_task.create_state());537 }538 }539 GGML_ASSERT(states.size() == id_tasks.size());540 queue_results.add_waiting_task_ids(id_tasks);541 queue_tasks.post(std::move(tasks), front);542}543 544bool server_response_reader::has_next() const {545 return !cancelled && received_count < id_tasks.size();546}547 548// return nullptr if should_stop() is true before receiving a result549// note: if one error is received, it will stop further processing and return error result550server_task_result_ptr server_response_reader::next(const std::function<bool()> & should_stop) {551 while (true) {552 server_task_result_ptr result = queue_results.recv_with_timeout(id_tasks, polling_interval_seconds);553 if (result == nullptr) {554 // timeout, check stop condition555 if (should_stop()) {556 return nullptr;557 }558 } else {559 if (result->is_error()) {560 stop(); // cancel remaining tasks561 SRV_DBG("%s", "received error result, stopping further processing\n");562 return result;563 }564 if (!states.empty()) {565 // update the generation state if needed566 const size_t idx = result->index;567 GGML_ASSERT(idx < states.size());568 result->update(states[idx]);569 }570 if (result->is_stop()) {571 received_count++;572 }573 return result;574 }575 }576 577 // should not reach here578}579 580server_response_reader::batch_response server_response_reader::wait_for_all(const std::function<bool()> & should_stop) {581 batch_response batch_res;582 batch_res.results.clear();583 batch_res.results.resize(id_tasks.size());584 while (has_next()) {585 auto res = next(should_stop);586 if (res == nullptr) {587 batch_res.is_terminated = true;588 return batch_res;589 }590 if (res->is_error()) {591 batch_res.error = std::move(res);592 return batch_res;593 }594 const size_t idx = res->index;595 GGML_ASSERT(idx < batch_res.results.size() && "index out of range");596 GGML_ASSERT(batch_res.results[idx] == nullptr && "duplicate result received");597 batch_res.results[idx] = std::move(res);598 }599 return batch_res;600}601 602void server_response_reader::stop() {603 queue_results.remove_waiting_task_ids(id_tasks);604 if (has_next() && !cancelled) {605 // if tasks is not finished yet, cancel them606 cancelled = true;607 std::vector<server_task> cancel_tasks;608 cancel_tasks.reserve(id_tasks.size());609 for (const auto & id_task : id_tasks) {610 SRV_WRN("cancel task, id_task = %d\n", id_task);611 server_task task(SERVER_TASK_TYPE_CANCEL);612 task.id_target = id_task;613 queue_results.remove_waiting_task_id(id_task);614 cancel_tasks.push_back(std::move(task));615 }616 // push to beginning of the queue, so it has highest priority617 queue_tasks.post(std::move(cancel_tasks), true);618 } else {619 SRV_DBG("%s", "all tasks already finished, no need to cancel\n");620 }621}622