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