ext_server.cpp 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380
  1. #include "ext_server.h"
  2. #include <atomic>
  3. // Necessary evil since the server types are not defined in a header
  4. #include "server.cpp"
  5. // Low level API access to verify GPU access
  6. #if defined(GGML_USE_CUBLAS)
  7. #if defined(GGML_USE_HIPBLAS)
  8. #include <hip/hip_runtime.h>
  9. #include <hipblas/hipblas.h>
  10. #include <hip/hip_fp16.h>
  11. #ifdef __HIP_PLATFORM_AMD__
  12. // for rocblas_initialize()
  13. #include "rocblas/rocblas.h"
  14. #endif // __HIP_PLATFORM_AMD__
  15. #define cudaGetDevice hipGetDevice
  16. #define cudaError_t hipError_t
  17. #define cudaSuccess hipSuccess
  18. #define cudaGetErrorString hipGetErrorString
  19. #else
  20. #include <cuda_runtime.h>
  21. #include <cublas_v2.h>
  22. #include <cuda_fp16.h>
  23. #endif // defined(GGML_USE_HIPBLAS)
  24. #endif // GGML_USE_CUBLAS
  25. // Expose the llama server as a callable extern "C" API
  26. llama_server_context *llama = NULL;
  27. std::thread ext_server_thread;
  28. bool shutting_down = false;
  29. std::atomic_int recv_counter;
  30. // RAII wrapper for tracking in-flight recv calls
  31. class atomicRecv {
  32. public:
  33. atomicRecv(std::atomic<int> &atomic) : atomic(atomic) {
  34. ++this->atomic;
  35. }
  36. ~atomicRecv() {
  37. --this->atomic;
  38. }
  39. private:
  40. std::atomic<int> &atomic;
  41. };
  42. void llama_server_init(ext_server_params *sparams, ext_server_resp_t *err) {
  43. recv_counter = 0;
  44. assert(err != NULL && sparams != NULL);
  45. log_set_target(stderr);
  46. if (!sparams->verbose_logging) {
  47. server_verbose = true;
  48. log_disable();
  49. }
  50. LOG_TEE("system info: %s\n", llama_print_system_info());
  51. err->id = 0;
  52. err->msg[0] = '\0';
  53. try {
  54. llama = new llama_server_context;
  55. gpt_params params;
  56. params.n_ctx = sparams->n_ctx;
  57. params.n_batch = sparams->n_batch;
  58. if (sparams->n_threads > 0) {
  59. params.n_threads = sparams->n_threads;
  60. }
  61. params.n_parallel = sparams->n_parallel;
  62. params.rope_freq_base = sparams->rope_freq_base;
  63. params.rope_freq_scale = sparams->rope_freq_scale;
  64. if (sparams->memory_f16) {
  65. params.cache_type_k = "f16";
  66. params.cache_type_v = "f16";
  67. } else {
  68. params.cache_type_k = "f32";
  69. params.cache_type_v = "f32";
  70. }
  71. params.n_gpu_layers = sparams->n_gpu_layers;
  72. params.main_gpu = sparams->main_gpu;
  73. params.use_mlock = sparams->use_mlock;
  74. params.use_mmap = sparams->use_mmap;
  75. params.numa = sparams->numa;
  76. params.embedding = sparams->embedding;
  77. if (sparams->model != NULL) {
  78. params.model = sparams->model;
  79. }
  80. if (sparams->lora_adapters != NULL) {
  81. for (ext_server_lora_adapter *la = sparams->lora_adapters; la != NULL;
  82. la = la->next) {
  83. params.lora_adapter.push_back(std::make_tuple(la->adapter, la->scale));
  84. }
  85. params.use_mmap = false;
  86. }
  87. if (sparams->mmproj != NULL) {
  88. params.mmproj = std::string(sparams->mmproj);
  89. }
  90. #if defined(GGML_USE_CUBLAS)
  91. // Before attempting to init the backend which will assert on error, verify the CUDA/ROCM GPU is accessible
  92. LOG_TEE("Performing pre-initialization of GPU\n");
  93. int id;
  94. cudaError_t cudaErr = cudaGetDevice(&id);
  95. if (cudaErr != cudaSuccess) {
  96. err->id = -1;
  97. snprintf(err->msg, err->msg_len, "Unable to init GPU: %s", cudaGetErrorString(cudaErr));
  98. return;
  99. }
  100. #endif
  101. llama_backend_init(params.numa);
  102. // load the model
  103. if (!llama->load_model(params)) {
  104. // TODO - consider modifying the logging logic or patching load_model so
  105. // we can capture more detailed error messages and pass them back to the
  106. // caller for better UX
  107. err->id = -1;
  108. snprintf(err->msg, err->msg_len, "error loading model %s",
  109. params.model.c_str());
  110. return;
  111. }
  112. llama->initialize();
  113. } catch (std::exception &e) {
  114. err->id = -1;
  115. snprintf(err->msg, err->msg_len, "exception %s", e.what());
  116. } catch (...) {
  117. err->id = -1;
  118. snprintf(err->msg, err->msg_len,
  119. "Unknown exception initializing llama server");
  120. }
  121. }
  122. void llama_server_start() {
  123. assert(llama != NULL);
  124. // TODO mutex to protect thread creation
  125. ext_server_thread = std::thread([&]() {
  126. try {
  127. LOG_TEE("llama server main loop starting\n");
  128. ggml_time_init();
  129. llama->queue_tasks.on_new_task(std::bind(
  130. &llama_server_context::process_single_task, llama, std::placeholders::_1));
  131. llama->queue_tasks.on_finish_multitask(std::bind(
  132. &llama_server_context::on_finish_multitask, llama, std::placeholders::_1));
  133. llama->queue_tasks.on_all_tasks_finished(std::bind(
  134. &llama_server_context::run_on_all_tasks_finished, llama));
  135. llama->queue_results.on_multitask_update(std::bind(
  136. &llama_server_queue::update_multitask,
  137. &llama->queue_tasks,
  138. std::placeholders::_1,
  139. std::placeholders::_2,
  140. std::placeholders::_3
  141. ));
  142. llama->queue_tasks.start_loop();
  143. } catch (std::exception &e) {
  144. LOG_TEE("caught exception in llama server main loop: %s\n", e.what());
  145. } catch (...) {
  146. LOG_TEE("caught unknown exception in llama server main loop\n");
  147. }
  148. LOG_TEE("\nllama server shutting down\n");
  149. llama_backend_free();
  150. });
  151. }
  152. void llama_server_stop() {
  153. assert(llama != NULL);
  154. // Shutdown any in-flight requests and block incoming requests.
  155. LOG_TEE("\ninitiating shutdown - draining remaining tasks...\n");
  156. shutting_down = true;
  157. while (recv_counter.load() > 0) {
  158. std::this_thread::sleep_for(std::chrono::milliseconds(50));
  159. }
  160. // This may take a while for any pending tasks to drain
  161. // TODO - consider a timeout to cancel tasks if it's taking too long
  162. llama->queue_tasks.terminate();
  163. ext_server_thread.join();
  164. delete llama;
  165. llama = NULL;
  166. LOG_TEE("llama server shutdown complete\n");
  167. }
  168. void llama_server_completion(const char *json_req, ext_server_resp_t *resp) {
  169. assert(llama != NULL && json_req != NULL && resp != NULL);
  170. resp->id = -1;
  171. resp->msg[0] = '\0';
  172. try {
  173. if (shutting_down) {
  174. throw std::runtime_error("server shutting down");
  175. }
  176. json data = json::parse(json_req);
  177. resp->id = llama->queue_tasks.get_new_id();
  178. llama->queue_results.add_waiting_task_id(resp->id);
  179. llama->request_completion(resp->id, data, false, false, -1);
  180. } catch (std::exception &e) {
  181. snprintf(resp->msg, resp->msg_len, "exception %s", e.what());
  182. } catch (...) {
  183. snprintf(resp->msg, resp->msg_len, "Unknown exception during completion");
  184. }
  185. }
  186. void llama_server_completion_next_result(const int task_id,
  187. ext_server_task_result_t *resp) {
  188. assert(llama != NULL && resp != NULL);
  189. std::string msg;
  190. resp->id = -1;
  191. resp->stop = false;
  192. resp->error = false;
  193. resp->json_resp = NULL;
  194. std::string result_json;
  195. try {
  196. atomicRecv ar(recv_counter);
  197. task_result result = llama->queue_results.recv(task_id);
  198. result_json =
  199. result.result_json.dump(-1, ' ', false, json::error_handler_t::replace);
  200. resp->id = result.id;
  201. resp->stop = result.stop;
  202. resp->error = result.error;
  203. if (result.error) {
  204. LOG_TEE("next result cancel on error\n");
  205. llama->request_cancel(task_id);
  206. LOG_TEE("next result removing waiting tak ID: %d\n", task_id);
  207. llama->queue_results.remove_waiting_task_id(task_id);
  208. } else if (result.stop) {
  209. LOG_TEE("next result cancel on stop\n");
  210. llama->request_cancel(task_id);
  211. LOG_TEE("next result removing waiting task ID: %d\n", task_id);
  212. llama->queue_results.remove_waiting_task_id(task_id);
  213. } else if (shutting_down) {
  214. LOG_TEE("aborting completion due to shutdown %d\n", task_id);
  215. llama->request_cancel(task_id);
  216. llama->queue_results.remove_waiting_task_id(task_id);
  217. resp->stop = true;
  218. }
  219. } catch (std::exception &e) {
  220. resp->error = true;
  221. resp->id = -1;
  222. result_json = "{\"error\":\"exception " + std::string(e.what()) + "\"}";
  223. LOG_TEE("llama server completion exception %s\n", e.what());
  224. } catch (...) {
  225. resp->error = true;
  226. resp->id = -1;
  227. result_json = "{\"error\":\"Unknown exception during completion\"}";
  228. LOG_TEE("llama server completion unknown exception\n");
  229. }
  230. const std::string::size_type size = result_json.size() + 1;
  231. resp->json_resp = new char[size];
  232. snprintf(resp->json_resp, size, "%s", result_json.c_str());
  233. }
  234. void llama_server_release_task_result(ext_server_task_result_t *result) {
  235. if (result == NULL || result->json_resp == NULL) {
  236. return;
  237. }
  238. delete[] result->json_resp;
  239. }
  240. void llama_server_completion_cancel(const int task_id, ext_server_resp_t *err) {
  241. assert(llama != NULL && err != NULL);
  242. err->id = 0;
  243. err->msg[0] = '\0';
  244. try {
  245. llama->request_cancel(task_id);
  246. llama->queue_results.remove_waiting_task_id(task_id);
  247. } catch (std::exception &e) {
  248. err->id = -1;
  249. snprintf(err->msg, err->msg_len, "exception %s", e.what());
  250. } catch (...) {
  251. err->id = -1;
  252. snprintf(err->msg, err->msg_len,
  253. "Unknown exception completion cancel in llama server");
  254. }
  255. }
  256. void llama_server_tokenize(const char *json_req, char **json_resp,
  257. ext_server_resp_t *err) {
  258. assert(llama != NULL && json_req != NULL && json_resp != NULL && err != NULL);
  259. *json_resp = NULL;
  260. err->id = 0;
  261. err->msg[0] = '\0';
  262. try {
  263. if (shutting_down) {
  264. throw std::runtime_error("server shutting down");
  265. }
  266. const json body = json::parse(json_req);
  267. std::vector<llama_token> tokens;
  268. if (body.count("content") != 0) {
  269. tokens = llama->tokenize(body["content"], false);
  270. }
  271. const json data = format_tokenizer_response(tokens);
  272. std::string result_json = data.dump();
  273. const std::string::size_type size = result_json.size() + 1;
  274. *json_resp = new char[size];
  275. snprintf(*json_resp, size, "%s", result_json.c_str());
  276. } catch (std::exception &e) {
  277. err->id = -1;
  278. snprintf(err->msg, err->msg_len, "exception %s", e.what());
  279. } catch (...) {
  280. err->id = -1;
  281. snprintf(err->msg, err->msg_len, "Unknown exception during tokenize");
  282. }
  283. }
  284. void llama_server_release_json_resp(char **json_resp) {
  285. if (json_resp == NULL || *json_resp == NULL) {
  286. return;
  287. }
  288. delete[] *json_resp;
  289. }
  290. void llama_server_detokenize(const char *json_req, char **json_resp,
  291. ext_server_resp_t *err) {
  292. assert(llama != NULL && json_req != NULL && json_resp != NULL && err != NULL);
  293. *json_resp = NULL;
  294. err->id = 0;
  295. err->msg[0] = '\0';
  296. try {
  297. if (shutting_down) {
  298. throw std::runtime_error("server shutting down");
  299. }
  300. const json body = json::parse(json_req);
  301. std::string content;
  302. if (body.count("tokens") != 0) {
  303. const std::vector<llama_token> tokens = body["tokens"];
  304. content = tokens_to_str(llama->ctx, tokens.cbegin(), tokens.cend());
  305. }
  306. const json data = format_detokenized_response(content);
  307. std::string result_json = data.dump();
  308. const std::string::size_type size = result_json.size() + 1;
  309. *json_resp = new char[size];
  310. snprintf(*json_resp, size, "%s", result_json.c_str());
  311. } catch (std::exception &e) {
  312. err->id = -1;
  313. snprintf(err->msg, err->msg_len, "exception %s", e.what());
  314. } catch (...) {
  315. err->id = -1;
  316. snprintf(err->msg, err->msg_len, "Unknown exception during detokenize");
  317. }
  318. }
  319. void llama_server_embedding(const char *json_req, char **json_resp,
  320. ext_server_resp_t *err) {
  321. assert(llama != NULL && json_req != NULL && json_resp != NULL && err != NULL);
  322. *json_resp = NULL;
  323. err->id = 0;
  324. err->msg[0] = '\0';
  325. try {
  326. if (shutting_down) {
  327. throw std::runtime_error("server shutting down");
  328. }
  329. const json body = json::parse(json_req);
  330. json prompt;
  331. if (body.count("content") != 0) {
  332. prompt = body["content"];
  333. } else {
  334. prompt = "";
  335. }
  336. const int task_id = llama->queue_tasks.get_new_id();
  337. llama->queue_results.add_waiting_task_id(task_id);
  338. llama->request_completion(task_id, {{"prompt", prompt}, {"n_predict", 0}}, false, true, -1);
  339. atomicRecv ar(recv_counter);
  340. task_result result = llama->queue_results.recv(task_id);
  341. std::string result_json = result.result_json.dump();
  342. const std::string::size_type size = result_json.size() + 1;
  343. *json_resp = new char[size];
  344. snprintf(*json_resp, size, "%s", result_json.c_str());
  345. llama->queue_results.remove_waiting_task_id(task_id);
  346. } catch (std::exception &e) {
  347. err->id = -1;
  348. snprintf(err->msg, err->msg_len, "exception %s", e.what());
  349. } catch (...) {
  350. err->id = -1;
  351. snprintf(err->msg, err->msg_len, "Unknown exception during embedding");
  352. }
  353. }