server.go 27 KB

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