server.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919
  1. package llm
  2. import (
  3. "bufio"
  4. "bytes"
  5. "context"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "log"
  11. "log/slog"
  12. "math/rand"
  13. "net"
  14. "net/http"
  15. "os"
  16. "os/exec"
  17. "path/filepath"
  18. "runtime"
  19. "strconv"
  20. "strings"
  21. "time"
  22. "golang.org/x/sync/semaphore"
  23. "github.com/ollama/ollama/api"
  24. "github.com/ollama/ollama/format"
  25. "github.com/ollama/ollama/gpu"
  26. )
  27. type LlamaServer interface {
  28. Ping(ctx context.Context) error
  29. WaitUntilRunning(ctx context.Context) error
  30. Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error
  31. Embedding(ctx context.Context, prompt string) ([]float64, error)
  32. Tokenize(ctx context.Context, content string) ([]int, error)
  33. Detokenize(ctx context.Context, tokens []int) (string, error)
  34. Close() error
  35. EstimatedVRAM() uint64
  36. }
  37. // llmServer is an instance of the llama.cpp server
  38. type llmServer struct {
  39. port int
  40. cmd *exec.Cmd
  41. done chan error // Channel to signal when the process exits
  42. status *StatusWriter
  43. options api.Options
  44. // TODO - this should be broken down by GPU
  45. estimatedVRAM uint64 // Estimated usage of VRAM by the loaded model
  46. sem *semaphore.Weighted
  47. }
  48. func LoadModel(model string) (*GGML, error) {
  49. if _, err := os.Stat(model); err != nil {
  50. return nil, err
  51. }
  52. f, err := os.Open(model)
  53. if err != nil {
  54. return nil, err
  55. }
  56. defer f.Close()
  57. ggml, _, err := DecodeGGML(f)
  58. return ggml, err
  59. }
  60. // NewLlamaServer will run a server for the given GPUs
  61. // The gpu list must be a single family.
  62. func NewLlamaServer(gpus gpu.GpuInfoList, model string, ggml *GGML, adapters, projectors []string, opts api.Options) (LlamaServer, error) {
  63. var err error
  64. if opts.NumCtx > int(ggml.KV().ContextLength()) {
  65. slog.Warn("requested context length is greater than the model's training context window size", "requested", opts.NumCtx, "training size", ggml.KV().ContextLength())
  66. }
  67. if opts.NumCtx < 4 {
  68. opts.NumCtx = 4
  69. }
  70. cpuRunner := ""
  71. var estimatedVRAM uint64
  72. var systemMemory uint64
  73. if (len(gpus) == 1 && gpus[0].Library == "cpu") || opts.NumGPU == 0 {
  74. // TODO evaluate system memory to see if we should block the load, or force an unload of another CPU runner
  75. cpuRunner = serverForCpu()
  76. } else {
  77. if gpus[0].Library == "metal" {
  78. memInfo, err := gpu.GetCPUMem()
  79. if err != nil {
  80. slog.Error("failed to lookup system memory", "error", err)
  81. } else {
  82. systemMemory = memInfo.TotalMemory
  83. slog.Debug("system memory", "total", format.HumanBytes2(systemMemory))
  84. }
  85. }
  86. var layers int
  87. layers, estimatedVRAM = EstimateGPULayers(gpus, ggml, projectors, opts)
  88. if gpus[0].Library == "metal" && estimatedVRAM > systemMemory {
  89. // disable partial offloading when model is greater than total system memory as this
  90. // can lead to locking up the system
  91. opts.NumGPU = 0
  92. } else if opts.NumGPU < 0 && layers > 0 && gpus[0].Library != "cpu" {
  93. opts.NumGPU = layers
  94. }
  95. }
  96. // Loop through potential servers
  97. finalErr := fmt.Errorf("no suitable llama servers found")
  98. if len(adapters) > 1 {
  99. return nil, errors.New("ollama supports only one lora adapter, but multiple were provided")
  100. }
  101. availableServers := availableServers()
  102. var servers []string
  103. if cpuRunner != "" {
  104. servers = []string{cpuRunner}
  105. } else {
  106. servers = serversForGpu(gpus[0]) // All GPUs in the list are matching Library and Variant
  107. }
  108. demandLib := strings.Trim(os.Getenv("OLLAMA_LLM_LIBRARY"), "\"' ")
  109. if demandLib != "" {
  110. serverPath := availableServers[demandLib]
  111. if serverPath == "" {
  112. slog.Info(fmt.Sprintf("Invalid OLLAMA_LLM_LIBRARY %s - not found", demandLib))
  113. } else {
  114. slog.Info("user override", "OLLAMA_LLM_LIBRARY", demandLib, "path", serverPath)
  115. servers = []string{demandLib}
  116. }
  117. }
  118. if len(servers) == 0 {
  119. return nil, fmt.Errorf("no servers found for %v", gpus)
  120. }
  121. params := []string{
  122. "--model", model,
  123. "--ctx-size", fmt.Sprintf("%d", opts.NumCtx),
  124. "--batch-size", fmt.Sprintf("%d", opts.NumBatch),
  125. "--embedding",
  126. }
  127. if debug := os.Getenv("OLLAMA_DEBUG"); debug != "" {
  128. params = append(params, "--log-format", "json")
  129. } else {
  130. params = append(params, "--log-disable")
  131. }
  132. if opts.NumGPU >= 0 {
  133. params = append(params, "--n-gpu-layers", fmt.Sprintf("%d", opts.NumGPU))
  134. }
  135. if debug := os.Getenv("OLLAMA_DEBUG"); debug != "" {
  136. params = append(params, "--verbose")
  137. }
  138. if opts.MainGPU > 0 {
  139. params = append(params, "--main-gpu", fmt.Sprintf("%d", opts.MainGPU))
  140. }
  141. if len(adapters) > 0 {
  142. // TODO: applying multiple adapters is not supported by the llama.cpp server yet
  143. params = append(params, "--lora", adapters[0])
  144. }
  145. if len(projectors) > 0 {
  146. // TODO: applying multiple projectors is not supported by the llama.cpp server yet
  147. params = append(params, "--mmproj", projectors[0])
  148. }
  149. if opts.NumThread > 0 {
  150. params = append(params, "--threads", fmt.Sprintf("%d", opts.NumThread))
  151. }
  152. if !opts.F16KV {
  153. params = append(params, "--memory-f32")
  154. }
  155. if opts.UseMLock {
  156. params = append(params, "--mlock")
  157. }
  158. if !opts.UseMMap {
  159. params = append(params, "--no-mmap")
  160. }
  161. if opts.UseNUMA {
  162. params = append(params, "--numa")
  163. }
  164. // "--cont-batching", // TODO - doesn't seem to have any noticeable perf change for multiple requests
  165. numParallel := 1
  166. if onp := os.Getenv("OLLAMA_NUM_PARALLEL"); onp != "" {
  167. numParallel, err = strconv.Atoi(onp)
  168. if err != nil || numParallel <= 0 {
  169. err = fmt.Errorf("invalid OLLAMA_NUM_PARALLEL=%s must be greater than zero - %w", onp, err)
  170. slog.Error("misconfiguration", "error", err)
  171. return nil, err
  172. }
  173. }
  174. params = append(params, "--parallel", fmt.Sprintf("%d", numParallel))
  175. for i := 0; i < len(servers); i++ {
  176. dir := availableServers[servers[i]]
  177. if dir == "" {
  178. // Shouldn't happen
  179. finalErr = fmt.Errorf("[%d] server %s not listed in available servers %v", i, servers[i], availableServers)
  180. slog.Error("sever list inconsistent", "error", finalErr)
  181. continue
  182. }
  183. // Find an availableServers port, retry on each iterration in case the failure was a port conflict race
  184. port := 0
  185. if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
  186. var l *net.TCPListener
  187. if l, err = net.ListenTCP("tcp", a); err == nil {
  188. port = l.Addr().(*net.TCPAddr).Port
  189. l.Close()
  190. }
  191. }
  192. if port == 0 {
  193. slog.Debug("ResolveTCPAddr failed ", "error", err)
  194. port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
  195. }
  196. finalParams := append(params, "--port", strconv.Itoa(port))
  197. pathEnv := "LD_LIBRARY_PATH"
  198. if runtime.GOOS == "windows" {
  199. pathEnv = "PATH"
  200. }
  201. // append the server directory to LD_LIBRARY_PATH/PATH
  202. libraryPaths := []string{dir}
  203. if libraryPath, ok := os.LookupEnv(pathEnv); ok {
  204. // Append our runner directory to the path
  205. // This will favor system libraries over our bundled library dependencies
  206. libraryPaths = append(filepath.SplitList(libraryPath), libraryPaths...)
  207. }
  208. // Note: we always put the dependency path first
  209. // since this was the exact version we verified for AMD GPUs
  210. // and we favor what the user had in their path
  211. if gpus[0].DependencyPath != "" {
  212. // TODO refine for multi-gpu support
  213. libraryPaths = append([]string{gpus[0].DependencyPath}, libraryPaths...)
  214. }
  215. server := filepath.Join(dir, "ollama_llama_server")
  216. if runtime.GOOS == "windows" {
  217. server = server + ".exe"
  218. }
  219. // Detect tmp cleaners wiping out the file
  220. _, err := os.Stat(server)
  221. if errors.Is(err, os.ErrNotExist) {
  222. slog.Warn("llama server disappeared, reinitializing payloads", "path", server, "error", err)
  223. err = Init()
  224. if err != nil {
  225. slog.Warn("failed to reinitialize payloads", "error", err)
  226. return nil, err
  227. }
  228. }
  229. s := &llmServer{
  230. port: port,
  231. cmd: exec.Command(server, finalParams...),
  232. status: NewStatusWriter(os.Stderr),
  233. options: opts,
  234. estimatedVRAM: estimatedVRAM,
  235. sem: semaphore.NewWeighted(int64(numParallel)),
  236. }
  237. libEnv := fmt.Sprintf("%s=%s", pathEnv, strings.Join(libraryPaths, string(filepath.ListSeparator)))
  238. s.cmd.Env = append(os.Environ(), libEnv)
  239. s.cmd.Stdout = os.Stdout
  240. s.cmd.Stderr = s.status
  241. // TODO - multiple GPU selection logic...
  242. key, val := gpu.GpuInfoList(gpus).GetVisibleDevicesEnv()
  243. if key != "" {
  244. s.cmd.Env = append(s.cmd.Env, key+"="+val)
  245. }
  246. slog.Info("starting llama server", "cmd", s.cmd.String())
  247. // Log at debug as the environment is inherited and might contain sensitive information
  248. slog.Debug("subprocess", "environment", s.cmd.Env)
  249. if err = s.cmd.Start(); err != nil {
  250. msg := ""
  251. if s.status != nil && s.status.LastErrMsg != "" {
  252. msg = s.status.LastErrMsg
  253. }
  254. err = fmt.Errorf("error starting the external llama server: %v %s", err, msg)
  255. finalErr = err
  256. continue
  257. }
  258. // reap subprocess when it exits
  259. go func() {
  260. // Exit status managed via getServerStatus
  261. _ = s.cmd.Wait()
  262. }()
  263. // TODO - make sure this is all wired up correctly
  264. // if err = s.WaitUntilRunning(); err != nil {
  265. // slog.Error("error starting llama server", "server", servers[i], "error", err)
  266. // s.Close()
  267. // finalErr = err
  268. // continue
  269. // }
  270. return s, nil
  271. }
  272. slog.Error("unable to load any llama server", "error", finalErr)
  273. return nil, finalErr
  274. }
  275. func projectorMemoryRequirements(filename string) uint64 {
  276. file, err := os.Open(filename)
  277. if err != nil {
  278. return 0
  279. }
  280. defer file.Close()
  281. ggml, _, err := DecodeGGML(file)
  282. if err != nil {
  283. return 0
  284. }
  285. var mem uint64
  286. for _, layer := range ggml.Tensors().Layers() {
  287. mem += layer.size()
  288. }
  289. return mem
  290. }
  291. type ServerStatus int
  292. const ( // iota is reset to 0
  293. ServerStatusReady ServerStatus = iota
  294. ServerStatusNoSlotsAvaialble
  295. ServerStatusLoadingModel
  296. ServerStatusNotResponding
  297. ServerStatusError
  298. )
  299. func (s ServerStatus) ToString() string {
  300. switch s {
  301. case ServerStatusReady:
  302. return "llm server ready"
  303. case ServerStatusNoSlotsAvaialble:
  304. return "llm busy - no slots available"
  305. case ServerStatusLoadingModel:
  306. return "llm server loading model"
  307. case ServerStatusNotResponding:
  308. return "llm server not responding"
  309. default:
  310. return "llm server error"
  311. }
  312. }
  313. type ServerStatusResp struct {
  314. Status string `json:"status"`
  315. SlotsIdle int `json:"slots_idle"`
  316. SlotsProcessing int `json:"slots_processing"`
  317. Error string `json:"error"`
  318. }
  319. func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
  320. // Fail fast if its exited
  321. if s.cmd.ProcessState != nil {
  322. msg := ""
  323. if s.status != nil && s.status.LastErrMsg != "" {
  324. msg = s.status.LastErrMsg
  325. }
  326. return ServerStatusError, fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
  327. }
  328. req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d/health", s.port), nil)
  329. if err != nil {
  330. return ServerStatusError, fmt.Errorf("error creating GET request: %v", err)
  331. }
  332. req.Header.Set("Content-Type", "application/json")
  333. resp, err := http.DefaultClient.Do(req)
  334. if err != nil {
  335. if errors.Is(err, context.DeadlineExceeded) {
  336. return ServerStatusNotResponding, fmt.Errorf("server not responding")
  337. }
  338. return ServerStatusError, fmt.Errorf("health resp: %w", err)
  339. }
  340. defer resp.Body.Close()
  341. body, err := io.ReadAll(resp.Body)
  342. if err != nil {
  343. return ServerStatusError, fmt.Errorf("read health request: %w", err)
  344. }
  345. var status ServerStatusResp
  346. if err := json.Unmarshal(body, &status); err != nil {
  347. return ServerStatusError, fmt.Errorf("health unmarshal encode response: %w", err)
  348. }
  349. switch status.Status {
  350. case "ok":
  351. return ServerStatusReady, nil
  352. case "no slot available":
  353. return ServerStatusNoSlotsAvaialble, nil
  354. case "loading model":
  355. return ServerStatusLoadingModel, nil
  356. default:
  357. return ServerStatusError, fmt.Errorf("server error: %+v", status)
  358. }
  359. }
  360. func (s *llmServer) Ping(ctx context.Context) error {
  361. _, err := s.getServerStatus(ctx)
  362. if err != nil {
  363. slog.Debug("server unhealthy", "error", err)
  364. return err
  365. }
  366. return nil
  367. }
  368. func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
  369. start := time.Now()
  370. // TODO we need to wire up a better way to detect hangs during model load and startup of the server
  371. expiresAt := time.Now().Add(10 * time.Minute) // be generous with timeout, large models can take a while to load
  372. ticker := time.NewTicker(50 * time.Millisecond)
  373. defer ticker.Stop()
  374. slog.Info("waiting for llama runner to start responding")
  375. var lastStatus ServerStatus = -1
  376. for {
  377. select {
  378. case <-ctx.Done():
  379. slog.Info("context expired before server started")
  380. return fmt.Errorf("timed out waiting for llama runner to start: %w", ctx.Err())
  381. case err := <-s.done:
  382. msg := ""
  383. if s.status != nil && s.status.LastErrMsg != "" {
  384. msg = s.status.LastErrMsg
  385. }
  386. return fmt.Errorf("llama runner process has terminated: %v %s", err, msg)
  387. case <-ticker.C:
  388. if time.Now().After(expiresAt) {
  389. // timeout
  390. msg := ""
  391. if s.status != nil && s.status.LastErrMsg != "" {
  392. msg = s.status.LastErrMsg
  393. }
  394. return fmt.Errorf("timed out waiting for llama runner to start: %s", msg)
  395. }
  396. if s.cmd.ProcessState != nil {
  397. msg := ""
  398. if s.status != nil && s.status.LastErrMsg != "" {
  399. msg = s.status.LastErrMsg
  400. }
  401. return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
  402. }
  403. c, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
  404. defer cancel()
  405. status, err := s.getServerStatus(c)
  406. if err != nil && lastStatus != status {
  407. slog.Debug("server not yet available", "error", err)
  408. lastStatus = status
  409. continue
  410. }
  411. switch status {
  412. case ServerStatusLoadingModel:
  413. // TODO - this state never seems to happen with the current server.cpp code (bug?)
  414. // it doesn't respond to the health endpoint until after the model is loaded
  415. slog.Debug("loading model")
  416. case ServerStatusReady:
  417. slog.Debug(fmt.Sprintf("llama runner started in %f seconds", time.Since(start).Seconds()))
  418. return nil
  419. }
  420. }
  421. }
  422. }
  423. const jsonGrammar = `
  424. root ::= object
  425. value ::= object | array | string | number | ("true" | "false" | "null") ws
  426. object ::=
  427. "{" ws (
  428. string ":" ws value
  429. ("," ws string ":" ws value)*
  430. )? "}" ws
  431. array ::=
  432. "[" ws (
  433. value
  434. ("," ws value)*
  435. )? "]" ws
  436. string ::=
  437. "\"" (
  438. [^"\\] |
  439. "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  440. )* "\"" ws
  441. number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws
  442. # Optional space: by convention, applied in this grammar after literal chars when allowed
  443. ws ::= ([ \t\n] ws)?
  444. `
  445. const maxBufferSize = 512 * format.KiloByte
  446. const maxRetries = 3
  447. type ImageData struct {
  448. Data []byte `json:"data"`
  449. ID int `json:"id"`
  450. }
  451. type completion struct {
  452. Content string `json:"content"`
  453. Model string `json:"model"`
  454. Prompt string `json:"prompt"`
  455. Stop bool `json:"stop"`
  456. Timings struct {
  457. PredictedN int `json:"predicted_n"`
  458. PredictedMS float64 `json:"predicted_ms"`
  459. PromptN int `json:"prompt_n"`
  460. PromptMS float64 `json:"prompt_ms"`
  461. }
  462. }
  463. type CompletionRequest struct {
  464. Prompt string
  465. Format string
  466. Images []ImageData
  467. Options api.Options
  468. }
  469. type CompletionResponse struct {
  470. Content string
  471. Done bool
  472. PromptEvalCount int
  473. PromptEvalDuration time.Duration
  474. EvalCount int
  475. EvalDuration time.Duration
  476. }
  477. func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
  478. if err := s.sem.Acquire(ctx, 1); err != nil {
  479. slog.Error("Failed to acquire semaphore", "error", err)
  480. return err
  481. }
  482. defer s.sem.Release(1)
  483. // only allow maximum 10 "context shifts" to avoid infinite generation
  484. if req.Options.NumPredict < 0 || req.Options.NumPredict > 10*s.options.NumCtx {
  485. req.Options.NumPredict = 10 * s.options.NumCtx
  486. slog.Debug("setting token limit to 10x num_ctx", "num_ctx", s.options.NumCtx, "num_predict", req.Options.NumPredict)
  487. }
  488. request := map[string]any{
  489. "prompt": req.Prompt,
  490. "stream": true,
  491. "n_predict": req.Options.NumPredict,
  492. "n_keep": req.Options.NumKeep,
  493. "main_gpu": req.Options.MainGPU,
  494. "temperature": req.Options.Temperature,
  495. "top_k": req.Options.TopK,
  496. "top_p": req.Options.TopP,
  497. "tfs_z": req.Options.TFSZ,
  498. "typical_p": req.Options.TypicalP,
  499. "repeat_last_n": req.Options.RepeatLastN,
  500. "repeat_penalty": req.Options.RepeatPenalty,
  501. "presence_penalty": req.Options.PresencePenalty,
  502. "frequency_penalty": req.Options.FrequencyPenalty,
  503. "mirostat": req.Options.Mirostat,
  504. "mirostat_tau": req.Options.MirostatTau,
  505. "mirostat_eta": req.Options.MirostatEta,
  506. "penalize_nl": req.Options.PenalizeNewline,
  507. "seed": req.Options.Seed,
  508. "stop": req.Options.Stop,
  509. "image_data": req.Images,
  510. "cache_prompt": true,
  511. }
  512. // Make sure the server is ready
  513. status, err := s.getServerStatus(ctx)
  514. if err != nil {
  515. return err
  516. } else if status != ServerStatusReady {
  517. return fmt.Errorf("unexpected server status: %s", status.ToString())
  518. }
  519. if req.Format == "json" {
  520. request["grammar"] = jsonGrammar
  521. if !strings.Contains(strings.ToLower(req.Prompt), "json") {
  522. slog.Warn("Prompt does not specify that the LLM should response in JSON, but JSON format is expected. For best results specify that JSON is expected in the system prompt.")
  523. }
  524. }
  525. retryDelay := 100 * time.Microsecond
  526. for retries := 0; retries < maxRetries; retries++ {
  527. if retries > 0 {
  528. time.Sleep(retryDelay) // wait before retrying
  529. retryDelay *= 2 // exponential backoff
  530. }
  531. // Handling JSON marshaling with special characters unescaped.
  532. buffer := &bytes.Buffer{}
  533. enc := json.NewEncoder(buffer)
  534. enc.SetEscapeHTML(false)
  535. if err := enc.Encode(request); err != nil {
  536. return fmt.Errorf("failed to marshal data: %v", err)
  537. }
  538. endpoint := fmt.Sprintf("http://127.0.0.1:%d/completion", s.port)
  539. req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, buffer)
  540. if err != nil {
  541. return fmt.Errorf("error creating POST request: %v", err)
  542. }
  543. req.Header.Set("Content-Type", "application/json")
  544. resp, err := http.DefaultClient.Do(req)
  545. if err != nil {
  546. return fmt.Errorf("POST predict: %v", err)
  547. }
  548. defer resp.Body.Close()
  549. if resp.StatusCode >= 400 {
  550. bodyBytes, err := io.ReadAll(resp.Body)
  551. if err != nil {
  552. return fmt.Errorf("failed reading llm error response: %w", err)
  553. }
  554. log.Printf("llm predict error: %s", bodyBytes)
  555. return fmt.Errorf("%s", bodyBytes)
  556. }
  557. scanner := bufio.NewScanner(resp.Body)
  558. buf := make([]byte, 0, maxBufferSize)
  559. scanner.Buffer(buf, maxBufferSize)
  560. retryNeeded := false
  561. // keep track of the last token generated, this is used to abort if the model starts looping
  562. var lastToken string
  563. var tokenRepeat int
  564. for scanner.Scan() {
  565. select {
  566. case <-ctx.Done():
  567. // This handles the request cancellation
  568. return ctx.Err()
  569. default:
  570. line := scanner.Bytes()
  571. if len(line) == 0 {
  572. continue
  573. }
  574. // try again on slot unavailable
  575. if bytes.Contains(line, []byte("slot unavailable")) {
  576. retryNeeded = true
  577. break
  578. }
  579. evt, ok := bytes.CutPrefix(line, []byte("data: "))
  580. if !ok {
  581. return fmt.Errorf("error parsing llm response stream: %s", line)
  582. }
  583. var c completion
  584. if err := json.Unmarshal(evt, &c); err != nil {
  585. return fmt.Errorf("error unmarshaling llm prediction response: %v", err)
  586. }
  587. switch {
  588. case strings.TrimSpace(c.Content) == lastToken:
  589. tokenRepeat++
  590. default:
  591. lastToken = strings.TrimSpace(c.Content)
  592. tokenRepeat = 0
  593. }
  594. // 30 picked as an arbitrary max token repeat limit, modify as needed
  595. if tokenRepeat > 30 {
  596. slog.Debug("prediction aborted, token repeat limit reached")
  597. return ctx.Err()
  598. }
  599. if c.Content != "" {
  600. fn(CompletionResponse{
  601. Content: c.Content,
  602. })
  603. }
  604. if c.Stop {
  605. fn(CompletionResponse{
  606. Done: true,
  607. PromptEvalCount: c.Timings.PromptN,
  608. PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
  609. EvalCount: c.Timings.PredictedN,
  610. EvalDuration: parseDurationMs(c.Timings.PredictedMS),
  611. })
  612. return nil
  613. }
  614. }
  615. }
  616. if err := scanner.Err(); err != nil {
  617. if strings.Contains(err.Error(), "unexpected EOF") {
  618. s.Close()
  619. msg := ""
  620. if s.status != nil && s.status.LastErrMsg != "" {
  621. msg = s.status.LastErrMsg
  622. }
  623. return fmt.Errorf("an unknown error was encountered while running the model %s", msg)
  624. }
  625. return fmt.Errorf("error reading llm response: %v", err)
  626. }
  627. if !retryNeeded {
  628. return nil // success
  629. }
  630. }
  631. // should never reach here ideally
  632. return fmt.Errorf("max retries exceeded")
  633. }
  634. type EmbeddingRequest struct {
  635. Content string `json:"content"`
  636. }
  637. type EmbeddingResponse struct {
  638. Embedding []float64 `json:"embedding"`
  639. }
  640. func (s *llmServer) Embedding(ctx context.Context, prompt string) ([]float64, error) {
  641. if err := s.sem.Acquire(ctx, 1); err != nil {
  642. slog.Error("Failed to acquire semaphore", "error", err)
  643. return nil, err
  644. }
  645. defer s.sem.Release(1)
  646. // Make sure the server is ready
  647. status, err := s.getServerStatus(ctx)
  648. if err != nil {
  649. return nil, err
  650. } else if status != ServerStatusReady {
  651. return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
  652. }
  653. data, err := json.Marshal(TokenizeRequest{Content: prompt})
  654. if err != nil {
  655. return nil, fmt.Errorf("error marshaling embed data: %w", err)
  656. }
  657. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/embedding", s.port), bytes.NewBuffer(data))
  658. if err != nil {
  659. return nil, fmt.Errorf("error creating embed request: %w", err)
  660. }
  661. req.Header.Set("Content-Type", "application/json")
  662. resp, err := http.DefaultClient.Do(req)
  663. if err != nil {
  664. return nil, fmt.Errorf("do embedding request: %w", err)
  665. }
  666. defer resp.Body.Close()
  667. body, err := io.ReadAll(resp.Body)
  668. if err != nil {
  669. return nil, fmt.Errorf("error reading embed response: %w", err)
  670. }
  671. if resp.StatusCode >= 400 {
  672. log.Printf("llm encode error: %s", body)
  673. return nil, fmt.Errorf("%s", body)
  674. }
  675. var embedding EmbeddingResponse
  676. if err := json.Unmarshal(body, &embedding); err != nil {
  677. return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
  678. }
  679. return embedding.Embedding, nil
  680. }
  681. type TokenizeRequest struct {
  682. Content string `json:"content"`
  683. }
  684. type TokenizeResponse struct {
  685. Tokens []int `json:"tokens"`
  686. }
  687. func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error) {
  688. // Make sure the server is ready
  689. status, err := s.getServerStatus(ctx)
  690. if err != nil {
  691. return nil, err
  692. } else if status != ServerStatusReady && status != ServerStatusNoSlotsAvaialble {
  693. return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
  694. }
  695. data, err := json.Marshal(TokenizeRequest{Content: content})
  696. if err != nil {
  697. return nil, fmt.Errorf("marshaling encode data: %w", err)
  698. }
  699. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/tokenize", s.port), bytes.NewBuffer(data))
  700. if err != nil {
  701. return nil, fmt.Errorf("encode request: %w", err)
  702. }
  703. req.Header.Set("Content-Type", "application/json")
  704. resp, err := http.DefaultClient.Do(req)
  705. if err != nil {
  706. return nil, fmt.Errorf("do encode request: %w", err)
  707. }
  708. defer resp.Body.Close()
  709. body, err := io.ReadAll(resp.Body)
  710. if err != nil {
  711. return nil, fmt.Errorf("read encode request: %w", err)
  712. }
  713. if resp.StatusCode >= 400 {
  714. log.Printf("llm encode error: %s", body)
  715. return nil, fmt.Errorf("%s", body)
  716. }
  717. var encoded TokenizeResponse
  718. if err := json.Unmarshal(body, &encoded); err != nil {
  719. return nil, fmt.Errorf("unmarshal encode response: %w", err)
  720. }
  721. return encoded.Tokens, nil
  722. }
  723. type DetokenizeRequest struct {
  724. Tokens []int `json:"tokens"`
  725. }
  726. type DetokenizeResponse struct {
  727. Content string `json:"content"`
  728. }
  729. func (s *llmServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
  730. // Make sure the server is ready
  731. status, err := s.getServerStatus(ctx)
  732. if err != nil {
  733. return "", err
  734. } else if status != ServerStatusReady && status != ServerStatusNoSlotsAvaialble {
  735. return "", fmt.Errorf("unexpected server status: %s", status.ToString())
  736. }
  737. data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
  738. if err != nil {
  739. return "", fmt.Errorf("marshaling decode data: %w", err)
  740. }
  741. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/detokenize", s.port), bytes.NewBuffer(data))
  742. if err != nil {
  743. return "", fmt.Errorf("decode request: %w", err)
  744. }
  745. req.Header.Set("Content-Type", "application/json")
  746. resp, err := http.DefaultClient.Do(req)
  747. if err != nil {
  748. return "", fmt.Errorf("do decode request: %w", err)
  749. }
  750. defer resp.Body.Close()
  751. body, err := io.ReadAll(resp.Body)
  752. if err != nil {
  753. return "", fmt.Errorf("read decode request: %w", err)
  754. }
  755. if resp.StatusCode >= 400 {
  756. log.Printf("llm decode error: %s", body)
  757. return "", fmt.Errorf("%s", body)
  758. }
  759. var decoded DetokenizeResponse
  760. if err := json.Unmarshal(body, &decoded); err != nil {
  761. return "", fmt.Errorf("unmarshal encode response: %w", err)
  762. }
  763. return decoded.Content, nil
  764. }
  765. func (s *llmServer) Close() error {
  766. if s.cmd != nil {
  767. slog.Debug("stopping llama server")
  768. return s.cmd.Process.Kill()
  769. }
  770. return nil
  771. }
  772. func (s *llmServer) EstimatedVRAM() uint64 {
  773. return s.estimatedVRAM
  774. }
  775. func parseDurationMs(ms float64) time.Duration {
  776. dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
  777. if err != nil {
  778. panic(err)
  779. }
  780. return dur
  781. }