ext_server.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428
  1. package llm
  2. /*
  3. #cgo CFLAGS: -I${SRCDIR}/llama.cpp/gguf -I${SRCDIR}/llama.cpp/gguf/common -I${SRCDIR}/llama.cpp/gguf/examples/server
  4. #cgo CFLAGS: -DNDEBUG -DLLAMA_SERVER_LIBRARY=1 -D_XOPEN_SOURCE=600 -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
  5. #cgo CFLAGS: -Wmissing-noreturn -Wall -Wextra -Wcast-qual -Wno-unused-function -Wno-array-bounds
  6. #cgo CPPFLAGS: -Ofast -Wall -Wextra -Wno-unused-function -Wno-unused-variable -Wno-deprecated-declarations -Wno-unused-but-set-variable
  7. #cgo darwin CFLAGS: -D_DARWIN_C_SOURCE
  8. #cgo darwin CPPFLAGS: -DGGML_USE_ACCELERATE
  9. #cgo darwin CPPFLAGS: -DGGML_USE_METAL -DGGML_METAL_NDEBUG
  10. #cgo darwin LDFLAGS: -lc++ -framework Accelerate
  11. #cgo darwin LDFLAGS: -framework Foundation -framework Metal -framework MetalKit -framework MetalPerformanceShaders
  12. #cgo darwin LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/metal/common/libcommon.a
  13. #cgo darwin LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/metal/examples/server/libext_server.a
  14. #cgo darwin LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/metal/libllama.a
  15. #cgo darwin LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/metal/libggml_static.a
  16. #cgo linux CFLAGS: -D_GNU_SOURCE
  17. #cgo linux windows CFLAGS: -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_MMV_Y=1 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_USE_CUBLAS
  18. #cgo linux LDFLAGS: -L/usr/local/cuda/targets/x86_64-linux/lib -L/usr/local/cuda/lib64 -L/usr/local/cuda/targets/x86_64-linux/lib/stubs
  19. #cgo linux LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/cpu/examples/server/libext_server.a
  20. #cgo linux LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/cpu/common/libcommon.a
  21. #cgo linux LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/cpu/libllama.a
  22. #cgo linux LDFLAGS: ${SRCDIR}/llama.cpp/gguf/build/cpu/libggml_static.a
  23. #cgo linux LDFLAGS: -lrt -lpthread -ldl -lstdc++ -lm
  24. #cgo windows LDFLAGS: -L${SRCDIR}/llama.cpp/gguf/build/wincuda/dist/bin
  25. #cgo windows LDFLAGS: -lext_server_shared -lpthread
  26. #include <stdlib.h>
  27. #include "server.h"
  28. */
  29. import "C"
  30. import (
  31. "bytes"
  32. "context"
  33. "encoding/json"
  34. "fmt"
  35. "log"
  36. "os"
  37. "runtime"
  38. "strings"
  39. "sync"
  40. "time"
  41. "unsafe"
  42. "github.com/jmorganca/ollama/api"
  43. "github.com/jmorganca/ollama/gpu"
  44. )
  45. func newExtServerResp(len C.size_t) C.ext_server_resp_t {
  46. var resp C.ext_server_resp_t
  47. resp.msg_len = len
  48. bytes := make([]byte, len)
  49. resp.msg = (*C.char)(C.CBytes(bytes))
  50. return resp
  51. }
  52. func freeExtServerResp(resp C.ext_server_resp_t) {
  53. if resp.msg_len == 0 {
  54. return
  55. }
  56. C.free(unsafe.Pointer(resp.msg))
  57. }
  58. func extServerResponseToErr(resp C.ext_server_resp_t) error {
  59. return fmt.Errorf(C.GoString(resp.msg))
  60. }
  61. type extServer interface {
  62. LLM
  63. llama_server_init(sparams *C.ext_server_params_t, err *C.ext_server_resp_t)
  64. llama_server_start()
  65. llama_server_stop()
  66. llama_server_completion(json_req *C.char, resp *C.ext_server_resp_t)
  67. llama_server_completion_next_result(task_id C.int, resp *C.ext_server_task_result_t)
  68. llama_server_completion_cancel(task_id C.int, err *C.ext_server_resp_t)
  69. llama_server_release_task_result(result *C.ext_server_task_result_t)
  70. llama_server_tokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t)
  71. llama_server_detokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t)
  72. llama_server_embedding(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t)
  73. llama_server_release_json_resp(json_resp **C.char)
  74. }
  75. type llamaExtServer struct {
  76. api.Options
  77. }
  78. // Note: current implementation does not support concurrent instantiations
  79. var mutex sync.Mutex
  80. func (llm *llamaExtServer) llama_server_init(sparams *C.ext_server_params_t, err *C.ext_server_resp_t) {
  81. C.llama_server_init(sparams, err)
  82. }
  83. func (llm *llamaExtServer) llama_server_start() {
  84. C.llama_server_start()
  85. }
  86. func (llm *llamaExtServer) llama_server_stop() {
  87. C.llama_server_stop()
  88. }
  89. func (llm *llamaExtServer) llama_server_completion(json_req *C.char, resp *C.ext_server_resp_t) {
  90. C.llama_server_completion(json_req, resp)
  91. }
  92. func (llm *llamaExtServer) llama_server_completion_next_result(task_id C.int, resp *C.ext_server_task_result_t) {
  93. C.llama_server_completion_next_result(task_id, resp)
  94. }
  95. func (llm *llamaExtServer) llama_server_completion_cancel(task_id C.int, err *C.ext_server_resp_t) {
  96. C.llama_server_completion_cancel(task_id, err)
  97. }
  98. func (llm *llamaExtServer) llama_server_release_task_result(result *C.ext_server_task_result_t) {
  99. C.llama_server_release_task_result(result)
  100. }
  101. func (llm *llamaExtServer) llama_server_tokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
  102. C.llama_server_tokenize(json_req, json_resp, err)
  103. }
  104. func (llm *llamaExtServer) llama_server_detokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
  105. C.llama_server_detokenize(json_req, json_resp, err)
  106. }
  107. func (llm *llamaExtServer) llama_server_embedding(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
  108. C.llama_server_embedding(json_req, json_resp, err)
  109. }
  110. func (llm *llamaExtServer) llama_server_release_json_resp(json_resp **C.char) {
  111. C.llama_server_release_json_resp(json_resp)
  112. }
  113. func newDefaultExtServer(model string, adapters, projectors []string, numLayers int64, opts api.Options) (extServer, error) {
  114. server := &llamaExtServer{opts}
  115. return newExtServer(server, model, adapters, projectors, numLayers, opts)
  116. }
  117. func newExtServer(server extServer, model string, adapters, projectors []string, numLayers int64, opts api.Options) (extServer, error) {
  118. if !mutex.TryLock() {
  119. log.Printf("concurrent llm servers not yet supported, waiting for prior server to complete")
  120. mutex.Lock()
  121. }
  122. fileInfo, err := os.Stat(model)
  123. if err != nil {
  124. return nil, err
  125. }
  126. var sparams C.ext_server_params_t
  127. sparams.model = C.CString(model)
  128. defer C.free(unsafe.Pointer(sparams.model))
  129. numGPU := gpu.NumGPU(numLayers, fileInfo.Size(), opts)
  130. sparams.embedding = true
  131. sparams.n_ctx = C.uint(opts.NumCtx)
  132. sparams.n_batch = C.uint(opts.NumBatch)
  133. sparams.n_gpu_layers = C.int(numGPU)
  134. sparams.main_gpu = C.int(opts.MainGPU)
  135. sparams.n_parallel = 1 // TODO - wire up concurrency
  136. // Always use the value encoded in the model
  137. sparams.rope_freq_base = 0.0
  138. sparams.rope_freq_scale = 0.0
  139. sparams.memory_f16 = C.bool(opts.F16KV)
  140. sparams.use_mlock = C.bool(opts.UseMLock)
  141. sparams.use_mmap = C.bool(opts.UseMMap)
  142. sparams.numa = C.bool(opts.UseNUMA)
  143. sparams.lora_adapters = nil
  144. for i := 0; i < len(adapters); i++ {
  145. la := (*C.ext_server_lora_adapter_t)(C.malloc(C.sizeof_ext_server_lora_adapter_t))
  146. defer C.free(unsafe.Pointer(la))
  147. la.adapter = C.CString(adapters[i])
  148. defer C.free(unsafe.Pointer(la.adapter))
  149. la.scale = C.float(1.0) // TODO expose scale/weights up through ollama UX
  150. la.next = nil
  151. if i == 0 {
  152. sparams.lora_adapters = la
  153. } else {
  154. tmp := sparams.lora_adapters
  155. for ; tmp.next != nil; tmp = tmp.next {
  156. }
  157. tmp.next = la
  158. }
  159. }
  160. if len(projectors) > 0 {
  161. // TODO: applying multiple projectors is not supported by the llama.cpp server yet
  162. sparams.mmproj = C.CString(projectors[0])
  163. defer C.free(unsafe.Pointer(sparams.mmproj))
  164. } else {
  165. sparams.mmproj = nil
  166. }
  167. if opts.NumThread > 0 {
  168. sparams.n_threads = C.uint(opts.NumThread)
  169. } else {
  170. sparams.n_threads = C.uint(runtime.NumCPU())
  171. }
  172. log.Printf("Initializing internal llama server")
  173. resp := newExtServerResp(128)
  174. defer freeExtServerResp(resp)
  175. server.llama_server_init(&sparams, &resp)
  176. if resp.id < 0 {
  177. return nil, extServerResponseToErr(resp)
  178. }
  179. log.Printf("Starting internal llama main loop")
  180. server.llama_server_start()
  181. return server, nil
  182. }
  183. func (llm *llamaExtServer) Predict(ctx context.Context, pred PredictOpts, fn func(PredictResult)) error {
  184. return predict(llm, llm.Options, ctx, pred, fn)
  185. }
  186. func predict(llm extServer, opts api.Options, ctx context.Context, predict PredictOpts, fn func(PredictResult)) error {
  187. resp := newExtServerResp(128)
  188. defer freeExtServerResp(resp)
  189. var imageData []ImageData
  190. if len(predict.Images) > 0 {
  191. for cnt, i := range predict.Images {
  192. imageData = append(imageData, ImageData{Data: i, ID: cnt})
  193. }
  194. }
  195. log.Printf("loaded %d images", len(imageData))
  196. request := map[string]any{
  197. "prompt": predict.Prompt,
  198. "stream": true,
  199. "n_predict": opts.NumPredict,
  200. "n_keep": opts.NumKeep,
  201. "temperature": opts.Temperature,
  202. "top_k": opts.TopK,
  203. "top_p": opts.TopP,
  204. "tfs_z": opts.TFSZ,
  205. "typical_p": opts.TypicalP,
  206. "repeat_last_n": opts.RepeatLastN,
  207. "repeat_penalty": opts.RepeatPenalty,
  208. "presence_penalty": opts.PresencePenalty,
  209. "frequency_penalty": opts.FrequencyPenalty,
  210. "mirostat": opts.Mirostat,
  211. "mirostat_tau": opts.MirostatTau,
  212. "mirostat_eta": opts.MirostatEta,
  213. "penalize_nl": opts.PenalizeNewline,
  214. "seed": opts.Seed,
  215. "stop": opts.Stop,
  216. "image_data": imageData,
  217. }
  218. if predict.Format == "json" {
  219. request["grammar"] = jsonGrammar
  220. }
  221. retryDelay := 100 * time.Microsecond
  222. for retries := 0; retries < maxRetries; retries++ {
  223. if retries > 0 {
  224. time.Sleep(retryDelay) // wait before retrying
  225. retryDelay *= 2 // exponential backoff
  226. }
  227. // Handling JSON marshaling with special characters unescaped.
  228. buffer := &bytes.Buffer{}
  229. enc := json.NewEncoder(buffer)
  230. enc.SetEscapeHTML(false)
  231. if err := enc.Encode(request); err != nil {
  232. return fmt.Errorf("failed to marshal data: %w", err)
  233. }
  234. req := C.CString(buffer.String())
  235. defer C.free(unsafe.Pointer(req))
  236. llm.llama_server_completion(req, &resp)
  237. if resp.id < 0 {
  238. return extServerResponseToErr(resp)
  239. }
  240. retryNeeded := false
  241. out:
  242. for {
  243. select {
  244. case <-ctx.Done():
  245. // This handles the request cancellation
  246. llm.llama_server_completion_cancel(resp.id, &resp)
  247. if resp.id < 0 {
  248. return extServerResponseToErr(resp)
  249. } else {
  250. return nil
  251. }
  252. default:
  253. var result C.ext_server_task_result_t
  254. llm.llama_server_completion_next_result(resp.id, &result)
  255. json_resp := C.GoString(result.json_resp)
  256. llm.llama_server_release_task_result(&result)
  257. var p prediction
  258. if err := json.Unmarshal([]byte(json_resp), &p); err != nil {
  259. llm.llama_server_completion_cancel(resp.id, &resp)
  260. if resp.id < 0 {
  261. return fmt.Errorf("error unmarshaling llm prediction response: %w and cancel %s", err, C.GoString(resp.msg))
  262. } else {
  263. return fmt.Errorf("error unmarshaling llm prediction response: %w", err)
  264. }
  265. }
  266. if bool(result.error) && strings.Contains(json_resp, "slot unavailable") {
  267. retryNeeded = true
  268. // task will already be canceled
  269. break out
  270. }
  271. if p.Content != "" {
  272. fn(PredictResult{
  273. Content: p.Content,
  274. })
  275. }
  276. if p.Stop {
  277. fn(PredictResult{
  278. Done: true,
  279. PromptEvalCount: p.Timings.PromptN,
  280. PromptEvalDuration: parseDurationMs(p.Timings.PromptMS),
  281. EvalCount: p.Timings.PredictedN,
  282. EvalDuration: parseDurationMs(p.Timings.PredictedMS),
  283. })
  284. return nil
  285. }
  286. }
  287. }
  288. if !retryNeeded {
  289. return nil // success
  290. }
  291. }
  292. // should never reach here ideally
  293. return fmt.Errorf("max retries exceeded")
  294. }
  295. func (llm *llamaExtServer) Encode(ctx context.Context, prompt string) ([]int, error) {
  296. return encode(llm, ctx, prompt)
  297. }
  298. func encode(llm extServer, ctx context.Context, prompt string) ([]int, error) {
  299. data, err := json.Marshal(TokenizeRequest{Content: prompt})
  300. if err != nil {
  301. return nil, fmt.Errorf("marshaling encode data: %w", err)
  302. }
  303. req := C.CString(string(data))
  304. defer C.free(unsafe.Pointer(req))
  305. var json_resp *C.char
  306. resp := newExtServerResp(128)
  307. defer freeExtServerResp(resp)
  308. llm.llama_server_tokenize(req, &json_resp, &resp)
  309. if resp.id < 0 {
  310. return nil, extServerResponseToErr(resp)
  311. }
  312. defer llm.llama_server_release_json_resp(&json_resp)
  313. var encoded TokenizeResponse
  314. if err2 := json.Unmarshal([]byte(C.GoString(json_resp)), &encoded); err2 != nil {
  315. return nil, fmt.Errorf("unmarshal encode response: %w", err2)
  316. }
  317. return encoded.Tokens, err
  318. }
  319. func (llm *llamaExtServer) Decode(ctx context.Context, tokens []int) (string, error) {
  320. return decode(llm, ctx, tokens)
  321. }
  322. func decode(llm extServer, ctx context.Context, tokens []int) (string, error) {
  323. if len(tokens) == 0 {
  324. return "", nil
  325. }
  326. data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
  327. if err != nil {
  328. return "", fmt.Errorf("marshaling decode data: %w", err)
  329. }
  330. req := C.CString(string(data))
  331. defer C.free(unsafe.Pointer(req))
  332. var json_resp *C.char
  333. resp := newExtServerResp(128)
  334. defer freeExtServerResp(resp)
  335. llm.llama_server_detokenize(req, &json_resp, &resp)
  336. if resp.id < 0 {
  337. return "", extServerResponseToErr(resp)
  338. }
  339. defer llm.llama_server_release_json_resp(&json_resp)
  340. var decoded DetokenizeResponse
  341. if err2 := json.Unmarshal([]byte(C.GoString(json_resp)), &decoded); err2 != nil {
  342. return "", fmt.Errorf("unmarshal encode response: %w", err2)
  343. }
  344. return decoded.Content, err
  345. }
  346. func (llm *llamaExtServer) Embedding(ctx context.Context, input string) ([]float64, error) {
  347. return embedding(llm, ctx, input)
  348. }
  349. func embedding(llm extServer, ctx context.Context, input string) ([]float64, error) {
  350. data, err := json.Marshal(TokenizeRequest{Content: input})
  351. if err != nil {
  352. return nil, fmt.Errorf("error marshaling embed data: %w", err)
  353. }
  354. req := C.CString(string(data))
  355. defer C.free(unsafe.Pointer(req))
  356. var json_resp *C.char
  357. resp := newExtServerResp(128)
  358. defer freeExtServerResp(resp)
  359. llm.llama_server_embedding(req, &json_resp, &resp)
  360. if resp.id < 0 {
  361. return nil, extServerResponseToErr(resp)
  362. }
  363. defer llm.llama_server_release_json_resp(&json_resp)
  364. var embedding EmbeddingResponse
  365. if err := json.Unmarshal([]byte(C.GoString(json_resp)), &embedding); err != nil {
  366. return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
  367. }
  368. return embedding.Embedding, nil
  369. }
  370. func (llm *llamaExtServer) Close() {
  371. close(llm)
  372. }
  373. func close(llm extServer) {
  374. llm.llama_server_stop()
  375. mutex.Unlock()
  376. }