vrr/prima.cpp
2
0
Fork 0
mirror of https://github.com/Lizonghang/prima.cpp.git synced 2025-09-05 19:49:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: block when free socket

Browse source
This commit is contained in:
DeEMO 2025-06-11 21:47:58 +08:00 • committed by DeEMO
parent 2039e3b0c1
commit d4618de991
2 changed files with 21 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
common/common.cpp
View file

@ -28,6 +28,7 @@
#include <unordered_set>
#include <vector>
#include <thread>
#include <atomic>
#if defined(__APPLE__) && defined(__MACH__)
#include <sys/types.h>
@ -1806,10 +1807,20 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
llama_update_context_with_rankworld(lctx, update_rank, update_n_world, worker_rank, n_worker);
if(node_type == NodeType::NODE_TYPE_FORWARDER){
//just foward
while (true) {
llama_forward_messages(lctx);
}
//just forward
std::atomic<bool> should_exit{false};
auto t = std::thread([lctx, &should_exit]() {
while(!should_exit) {
llama_forward_messages(lctx);
}
});
char * stop_signal = nullptr;
llama_free_sockets(lctx, &stop_signal); // this will block until receive stop signal
should_exit = true;
t.join();
exit(0);
}
// update n_layer_window and n_gpu_layers

8
src/llama.cpp
View file

@ -20776,9 +20776,13 @@ int llama_rebuild_topo(llama_context * ctx,
int llama_forward_messages(llama_context *ctx) {
zmq::message_t message;
int more = true;
int timeout_ms = 10;
ctx->recv_socket->setsockopt(ZMQ_RCVTIMEO, &timeout_ms, sizeof(timeout_ms));
while (more) {
ctx->recv_socket->recv(message, zmq::recv_flags::none);
auto recv_result = ctx->recv_socket->recv(message, zmq::recv_flags::none);
if (!recv_result) {
return -1;
}
size_t more_size = sizeof(more);
ctx->recv_socket->getsockopt(ZMQ_RCVMORE, &more, &more_size);