server.go 26 KB

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