vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

server : add more clean up when cancel_tasks is called

Browse source
This commit is contained in:
Xuan Son Nguyen 2025-01-22 00:04:09 +01:00
parent 6171c9d258
commit f4dab52a24
1 changed files with 36 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

42
examples/server/server.cpp
View file

@ -1428,6 +1428,10 @@ struct server_queue {
} else { } else {
queue_tasks.push_back(std::move(task)); queue_tasks.push_back(std::move(task));
} }
// if this is cancel task make sure to clean up pending tasks
if (task.type == SERVER_TASK_TYPE_CANCEL) {
cleanup_pending_task(task.id_target);
}
condition_tasks.notify_one(); condition_tasks.notify_one();
return task.id; return task.id;
} }
@ -1445,6 +1449,10 @@ struct server_queue {
} else { } else {
queue_tasks.push_back(std::move(task)); queue_tasks.push_back(std::move(task));
} }
// if this is cancel task make sure to clean up pending tasks
if (task.type == SERVER_TASK_TYPE_CANCEL) {
cleanup_pending_task(task.id_target);
}
} }
condition_tasks.notify_one(); condition_tasks.notify_one();
return 0; return 0;
@ -1539,6 +1547,23 @@ struct server_queue {
} }
} }
} }
private:
void cleanup_pending_task(int id_task) {
// no need lock because this is called exclusively by post()
for (size_t i = 0; i < queue_tasks.size(); i++) {
if (queue_tasks[i].id == id_task) {
queue_tasks.erase(queue_tasks.begin() + i);
break;
}
}
for (size_t i = 0; i < queue_tasks_deferred.size(); i++) {
if (queue_tasks_deferred[i].id == id_task) {
queue_tasks_deferred.erase(queue_tasks_deferred.begin() + i);
break;
}
}
}
}; };
struct server_response { struct server_response {
@ -1574,6 +1599,13 @@ struct server_response {
std::unique_lock<std::mutex> lock(mutex_results); std::unique_lock<std::mutex> lock(mutex_results);
waiting_task_ids.erase(id_task); waiting_task_ids.erase(id_task);
// make sure to clean up all pending results
for (size_t i = 0; i < queue_results.size(); i++) {
if (queue_results[i]->id == id_task) {
queue_results.erase(queue_results.begin() + i);
i--;
}
}
} }
void remove_waiting_task_ids(const std::unordered_set<int> & id_tasks) { void remove_waiting_task_ids(const std::unordered_set<int> & id_tasks) {
@ -1593,7 +1625,7 @@ struct server_response {
return !queue_results.empty(); return !queue_results.empty();
}); });
for (int i = 0; i < (int) queue_results.size(); i++) { for (size_t i = 0; i < queue_results.size(); i++) {
if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) { if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) {
server_task_result_ptr res = std::move(queue_results[i]); server_task_result_ptr res = std::move(queue_results[i]);
queue_results.erase(queue_results.begin() + i); queue_results.erase(queue_results.begin() + i);
@ -1610,10 +1642,8 @@ struct server_response {
server_task_result_ptr recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout) { server_task_result_ptr recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout) {
while (true) { while (true) {
std::unique_lock<std::mutex> lock(mutex_results); std::unique_lock<std::mutex> lock(mutex_results);
bool cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout), [&]{ std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout));
return !queue_results.empty(); if (cr_res == std::cv_status::timeout) {
});
if (!cr_res) {
return nullptr; return nullptr;
} }
@ -2368,8 +2398,8 @@ struct server_context {
server_task task(SERVER_TASK_TYPE_CANCEL); server_task task(SERVER_TASK_TYPE_CANCEL);
task.id_target = id_task; task.id_target = id_task;
cancel_tasks.push_back(task);
queue_results.remove_waiting_task_id(id_task); queue_results.remove_waiting_task_id(id_task);
cancel_tasks.push_back(task);
} }
// push to beginning of the queue, so it has highest priority // push to beginning of the queue, so it has highest priority
queue_tasks.post(cancel_tasks, true); queue_tasks.post(cancel_tasks, true);