server.go 32.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"log/slog"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
22
	"sync"
23
24
	"time"

Daniel Hiltgen's avatar
Daniel Hiltgen committed
25
26
	"golang.org/x/sync/semaphore"

27
	"github.com/ollama/ollama/api"
28
	"github.com/ollama/ollama/build"
29
	"github.com/ollama/ollama/discover"
30
	"github.com/ollama/ollama/envconfig"
31
	"github.com/ollama/ollama/format"
32
	"github.com/ollama/ollama/llama"
33
	"github.com/ollama/ollama/runners"
34
35
)

Daniel Hiltgen's avatar
Daniel Hiltgen committed
36
37
38
39
type LlamaServer interface {
	Ping(ctx context.Context) error
	WaitUntilRunning(ctx context.Context) error
	Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error
40
	Embedding(ctx context.Context, input string) ([]float32, error)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
41
42
43
	Tokenize(ctx context.Context, content string) ([]int, error)
	Detokenize(ctx context.Context, tokens []int) (string, error)
	Close() error
44
	EstimatedVRAM() uint64 // Total VRAM across all GPUs
45
	EstimatedTotal() uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
46
	EstimatedVRAMByGPU(gpuID string) uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
47
48
49
50
}

// llmServer is an instance of the llama.cpp server
type llmServer struct {
51
52
53
54
55
56
	port        int
	cmd         *exec.Cmd
	done        chan error // Channel to signal when the process exits
	status      *StatusWriter
	options     api.Options
	numParallel int
57
58
59
	modelPath   string
	modelLock   sync.Mutex   // Temporary until we switch fully to Go server
	model       *llama.Model // If non-nil, the runner is a new Go server
Daniel Hiltgen's avatar
Daniel Hiltgen committed
60

61
62
63
	estimate    MemoryEstimate
	totalLayers uint64
	// gpuCount     int
64
65
	gpus         discover.GpuInfoList // Recorded just before the model loaded, free space will be incorrect
	loadDuration time.Duration        // Record how long it took the model to load
66
	loadProgress float32
Daniel Hiltgen's avatar
Daniel Hiltgen committed
67
68

	sem *semaphore.Weighted
69
70
}

71
72
73
74
75
76
// LoadModel will load a model from disk. The model must be in the GGML format.
//
// It collects array values for arrays with a size less than or equal to
// maxArraySize. If maxArraySize is 0, the default value of 1024 is used. If
// the maxArraySize is negative, all arrays are collected.
func LoadModel(model string, maxArraySize int) (*GGML, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
77
78
79
80
	if _, err := os.Stat(model); err != nil {
		return nil, err
	}

81
82
83
84
85
86
	f, err := os.Open(model)
	if err != nil {
		return nil, err
	}
	defer f.Close()

87
	ggml, _, err := DecodeGGML(f, maxArraySize)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
88
89
	return ggml, err
}
90

Daniel Hiltgen's avatar
Daniel Hiltgen committed
91
92
// NewLlamaServer will run a server for the given GPUs
// The gpu list must be a single family.
93
func NewLlamaServer(gpus discover.GpuInfoList, model string, ggml *GGML, adapters, projectors []string, opts api.Options, numParallel int) (LlamaServer, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
94
	var err error
95
	var cpuRunner string
96
	var estimate MemoryEstimate
97
98
	var systemTotalMemory uint64
	var systemFreeMemory uint64
99
	var systemSwapFreeMemory uint64
100

101
	systemInfo := discover.GetSystemInfo()
102
103
104
105
	systemTotalMemory = systemInfo.System.TotalMemory
	systemFreeMemory = systemInfo.System.FreeMemory
	systemSwapFreeMemory = systemInfo.System.FreeSwap
	slog.Info("system memory", "total", format.HumanBytes2(systemTotalMemory), "free", format.HumanBytes2(systemFreeMemory), "free_swap", format.HumanBytes2(systemSwapFreeMemory))
106

107
108
	// If the user wants zero GPU layers, reset the gpu list to be CPU/system ram info
	if opts.NumGPU == 0 {
109
		gpus = discover.GetCPUInfo()
110
111
	}
	if len(gpus) == 1 && gpus[0].Library == "cpu" {
112
		cpuRunner = runners.ServerForCpu()
113
		estimate = EstimateGPULayers(gpus, ggml, projectors, opts)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
114
	} else {
115
		estimate = EstimateGPULayers(gpus, ggml, projectors, opts)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
116

Michael Yang's avatar
Michael Yang committed
117
		switch {
118
		case gpus[0].Library == "metal" && estimate.VRAMSize > systemTotalMemory:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
119
120
121
			// disable partial offloading when model is greater than total system memory as this
			// can lead to locking up the system
			opts.NumGPU = 0
122
		case gpus[0].Library != "metal" && estimate.Layers == 0:
123
			// Don't bother loading into the GPU if no layers can fit
124
			cpuRunner = runners.ServerForCpu()
125
			gpus = discover.GetCPUInfo()
126
127
		case opts.NumGPU < 0 && estimate.Layers > 0 && gpus[0].Library != "cpu":
			opts.NumGPU = estimate.Layers
128
129
130
		}
	}

131
132
133
	// On linux and windows, over-allocating CPU memory will almost always result in an error
	// Darwin has fully dynamic swap so has no direct concept of free swap space
	if runtime.GOOS != "darwin" {
134
		systemMemoryRequired := estimate.TotalSize - estimate.VRAMSize
135
		available := systemFreeMemory + systemSwapFreeMemory
136
137
138
		if systemMemoryRequired > available {
			slog.Warn("model request too large for system", "requested", format.HumanBytes2(systemMemoryRequired), "available", available, "total", format.HumanBytes2(systemTotalMemory), "free", format.HumanBytes2(systemFreeMemory), "swap", format.HumanBytes2(systemSwapFreeMemory))
			return nil, fmt.Errorf("model requires more system memory (%s) than is available (%s)", format.HumanBytes2(systemMemoryRequired), format.HumanBytes2(available))
139
140
141
		}
	}

142
143
	estimate.log()

Daniel Hiltgen's avatar
Daniel Hiltgen committed
144
	// Loop through potential servers
Michael Yang's avatar
Michael Yang committed
145
	finalErr := errors.New("no suitable llama servers found")
146

147
148
149
150
151
152
	rDir, err := runners.Refresh(build.EmbedFS)
	if err != nil {
		return nil, err
	}

	availableServers := runners.GetAvailableServers(rDir)
153
	if len(availableServers) == 0 {
154
		return nil, finalErr
155
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
156
157
158
159
	var servers []string
	if cpuRunner != "" {
		servers = []string{cpuRunner}
	} else {
160
		servers = runners.ServersForGpu(gpus[0]) // All GPUs in the list are matching Library and Variant
Daniel Hiltgen's avatar
Daniel Hiltgen committed
161
	}
Michael Yang's avatar
string  
Michael Yang committed
162
	demandLib := envconfig.LLMLibrary()
163
164
165
166
167
168
169
	if demandLib != "" {
		serverPath := availableServers[demandLib]
		if serverPath == "" {
			slog.Info(fmt.Sprintf("Invalid OLLAMA_LLM_LIBRARY %s - not found", demandLib))
		} else {
			slog.Info("user override", "OLLAMA_LLM_LIBRARY", demandLib, "path", serverPath)
			servers = []string{demandLib}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
170
171
172
173
			if strings.HasPrefix(demandLib, "cpu") {
				// Omit the GPU flag to silence the warning
				opts.NumGPU = -1
			}
174
175
176
177
		}
	}

	if len(servers) == 0 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
178
		return nil, fmt.Errorf("no servers found for %v", gpus)
179
180
181
182
	}

	params := []string{
		"--model", model,
Michael Yang's avatar
lint  
Michael Yang committed
183
184
		"--ctx-size", strconv.Itoa(opts.NumCtx),
		"--batch-size", strconv.Itoa(opts.NumBatch),
185
	}
Michael Yang's avatar
Michael Yang committed
186

Michael Yang's avatar
Michael Yang committed
187
	if opts.NumGPU >= 0 {
Michael Yang's avatar
lint  
Michael Yang committed
188
		params = append(params, "--n-gpu-layers", strconv.Itoa(opts.NumGPU))
189
190
	}

Michael Yang's avatar
Michael Yang committed
191
	if envconfig.Debug() {
192
193
194
195
		params = append(params, "--verbose")
	}

	if opts.MainGPU > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
196
		params = append(params, "--main-gpu", strconv.Itoa(opts.MainGPU))
197
198
199
	}

	if len(adapters) > 0 {
200
201
202
		for _, adapter := range adapters {
			params = append(params, "--lora", adapter)
		}
203
204
205
206
207
208
209
	}

	if len(projectors) > 0 {
		// TODO: applying multiple projectors is not supported by the llama.cpp server yet
		params = append(params, "--mmproj", projectors[0])
	}

210
	defaultThreads := systemInfo.GetOptimalThreadCount()
211
	if opts.NumThread > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
212
		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
213
214
	} else if defaultThreads > 0 {
		params = append(params, "--threads", strconv.Itoa(defaultThreads))
215
216
	}

Michael Yang's avatar
bool  
Michael Yang committed
217
	flashAttnEnabled := envconfig.FlashAttention()
Sam's avatar
Sam committed
218
219

	for _, g := range gpus {
220
		// only cuda (compute capability 7+) and metal support flash attention
Sam's avatar
Sam committed
221
		if g.Library != "metal" && (g.Library != "cuda" || g.DriverMajor < 7) {
222
			flashAttnEnabled = false
Sam's avatar
Sam committed
223
		}
224
225
226
227
228

		// mmap has issues with partial offloading on metal
		if g.Library == "metal" &&
			uint64(opts.NumGPU) > 0 &&
			uint64(opts.NumGPU) < ggml.KV().BlockCount()+1 {
229
230
			opts.UseMMap = new(bool)
			*opts.UseMMap = false
231
		}
Sam's avatar
Sam committed
232
	}
233

234
	if flashAttnEnabled {
Sam's avatar
Sam committed
235
236
237
		params = append(params, "--flash-attn")
	}

238
	// Windows CUDA should not use mmap for best performance
239
	// Linux  with a model larger than free space, mmap leads to thrashing
Daniel Hiltgen's avatar
Daniel Hiltgen committed
240
	// For CPU loads we want the memory to be allocated, not FS cache
241
242
243
244
	if (runtime.GOOS == "windows" && gpus[0].Library == "cuda" && opts.UseMMap == nil) ||
		(runtime.GOOS == "linux" && systemFreeMemory < estimate.TotalSize && opts.UseMMap == nil) ||
		(gpus[0].Library == "cpu" && opts.UseMMap == nil) ||
		(opts.UseMMap != nil && !*opts.UseMMap) {
245
246
247
248
249
250
251
		params = append(params, "--no-mmap")
	}

	if opts.UseMLock {
		params = append(params, "--mlock")
	}

252
	// TODO - NUMA support currently doesn't work properly
253

Michael Yang's avatar
lint  
Michael Yang committed
254
	params = append(params, "--parallel", strconv.Itoa(numParallel))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
255

256
257
258
259
	if estimate.TensorSplit != "" {
		params = append(params, "--tensor-split", estimate.TensorSplit)
	}

260
261
262
263
	if envconfig.MultiUserCache() {
		params = append(params, "--multiuser-cache")
	}

264
	for i := range servers {
265
		dir := availableServers[servers[i]]
266
267
268
		if dir == "" {
			// Shouldn't happen
			finalErr = fmt.Errorf("[%d] server %s not listed in available servers %v", i, servers[i], availableServers)
Michael Yang's avatar
Michael Yang committed
269
			slog.Error("server list inconsistent", "error", finalErr)
270
271
			continue
		}
272

Daniel Hiltgen's avatar
Daniel Hiltgen committed
273
		if strings.HasPrefix(servers[i], "cpu") {
274
			gpus = discover.GetCPUInfo()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
275
276
		}

277
		// Find an availableServers  port, retry on each iteration in case the failure was a port conflict race
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
		port := 0
		if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
			var l *net.TCPListener
			if l, err = net.ListenTCP("tcp", a); err == nil {
				port = l.Addr().(*net.TCPAddr).Port
				l.Close()
			}
		}
		if port == 0 {
			slog.Debug("ResolveTCPAddr failed ", "error", err)
			port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
		}
		finalParams := append(params, "--port", strconv.Itoa(port))

		pathEnv := "LD_LIBRARY_PATH"
		if runtime.GOOS == "windows" {
			pathEnv = "PATH"
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
296
297
		// Start with the server directory for the LD_LIBRARY_PATH/PATH
		libraryPaths := []string{dir}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
298

299
		if libraryPath, ok := os.LookupEnv(pathEnv); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
300
			// favor our bundled library dependencies over system libraries
Daniel Hiltgen's avatar
Daniel Hiltgen committed
301
			libraryPaths = append(libraryPaths, filepath.SplitList(libraryPath)...)
302
303
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
304
		// Note: we always put the dependency path first
Daniel Hiltgen's avatar
Daniel Hiltgen committed
305
		// since this was the exact version we compiled/linked against
306
		if gpus[0].DependencyPath != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
307
			// assume gpus from the same library have the same dependency path
308
			libraryPaths = append(gpus[0].DependencyPath, libraryPaths...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
309
310
		}

311
312
		server := filepath.Join(dir, "ollama_llama_server")
		if runtime.GOOS == "windows" {
Michael Yang's avatar
Michael Yang committed
313
			server += ".exe"
314
315
		}

316
317
318
319
		// Detect tmp cleaners wiping out the file
		_, err := os.Stat(server)
		if errors.Is(err, os.ErrNotExist) {
			slog.Warn("llama server disappeared, reinitializing payloads", "path", server, "error", err)
320
			_, err = runners.Refresh(build.EmbedFS)
321
322
323
324
325
326
			if err != nil {
				slog.Warn("failed to reinitialize payloads", "error", err)
				return nil, err
			}
		}

327
		// TODO - once fully switched to the Go runner, load the model here for tokenize/detokenize cgo access
Daniel Hiltgen's avatar
Daniel Hiltgen committed
328
		s := &llmServer{
329
330
331
332
			port:        port,
			cmd:         exec.Command(server, finalParams...),
			status:      NewStatusWriter(os.Stderr),
			options:     opts,
333
			modelPath:   model,
334
			estimate:    estimate,
335
			numParallel: numParallel,
336
337
			sem:         semaphore.NewWeighted(int64(numParallel)),
			totalLayers: ggml.KV().BlockCount() + 1,
338
			gpus:        gpus,
339
			done:        make(chan error, 1),
340
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
341

342
		s.cmd.Env = os.Environ()
343
344
		s.cmd.Stdout = os.Stdout
		s.cmd.Stderr = s.status
345
		s.cmd.SysProcAttr = LlamaServerSysProcAttr
346

Daniel Hiltgen's avatar
Daniel Hiltgen committed
347
348
349
350
		envWorkarounds := [][2]string{}
		for _, gpu := range gpus {
			envWorkarounds = append(envWorkarounds, gpu.EnvWorkarounds...)
		}
Michael Yang's avatar
lint  
Michael Yang committed
351
		visibleDevicesEnv, visibleDevicesEnvVal := gpus.GetVisibleDevicesEnv()
352
353
354
355
356
357
358
359
360
361
362
363
364
		pathEnvVal := strings.Join(libraryPaths, string(filepath.ListSeparator))

		// Update or add the path and visible devices variable with our adjusted version
		pathNeeded := true
		devicesNeeded := visibleDevicesEnv != ""
		for i := range s.cmd.Env {
			cmp := strings.SplitN(s.cmd.Env[i], "=", 2)
			if strings.EqualFold(cmp[0], pathEnv) {
				s.cmd.Env[i] = pathEnv + "=" + pathEnvVal
				pathNeeded = false
			} else if devicesNeeded && strings.EqualFold(cmp[0], visibleDevicesEnv) {
				s.cmd.Env[i] = visibleDevicesEnv + "=" + visibleDevicesEnvVal
				devicesNeeded = false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
365
366
367
368
369
370
			} else if len(envWorkarounds) != 0 {
				for _, kv := range envWorkarounds {
					if strings.EqualFold(cmp[0], kv[0]) {
						s.cmd.Env[i] = kv[0] + "=" + kv[1]
					}
				}
371
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
372
		}
373
374
		if pathNeeded {
			s.cmd.Env = append(s.cmd.Env, pathEnv+"="+pathEnvVal)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
375
		}
376
377
		if devicesNeeded {
			s.cmd.Env = append(s.cmd.Env, visibleDevicesEnv+"="+visibleDevicesEnvVal)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
378
379
		}

380
		slog.Info("starting llama server", "cmd", s.cmd.String())
Michael Yang's avatar
Michael Yang committed
381
		if envconfig.Debug() {
382
383
384
			filteredEnv := []string{}
			for _, ev := range s.cmd.Env {
				if strings.HasPrefix(ev, "CUDA_") ||
Daniel Hiltgen's avatar
Daniel Hiltgen committed
385
					strings.HasPrefix(ev, "ROCR_") ||
386
387
					strings.HasPrefix(ev, "ROCM_") ||
					strings.HasPrefix(ev, "HIP_") ||
Daniel Hiltgen's avatar
Daniel Hiltgen committed
388
					strings.HasPrefix(ev, "GPU_") ||
389
390
391
392
393
394
395
396
397
398
					strings.HasPrefix(ev, "HSA_") ||
					strings.HasPrefix(ev, "GGML_") ||
					strings.HasPrefix(ev, "PATH=") ||
					strings.HasPrefix(ev, "LD_LIBRARY_PATH=") {
					filteredEnv = append(filteredEnv, ev)
				}
			}
			// Log at debug as the environment is inherited and might contain sensitive information
			slog.Debug("subprocess", "environment", filteredEnv)
		}
399
400

		if err = s.cmd.Start(); err != nil {
401
			// Detect permission denied and augment the message about noexec
402
403
404
405
			if errors.Is(err, os.ErrPermission) {
				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)
				continue
			}
406
407
408
409
410
411
412
413
414
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
			err = fmt.Errorf("error starting the external llama server: %v %s", err, msg)
			finalErr = err
			continue
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
415
416
		// reap subprocess when it exits
		go func() {
417
418
419
420
421
422
423
			err := s.cmd.Wait()
			// Favor a more detailed message over the process exit status
			if err != nil && s.status != nil && s.status.LastErrMsg != "" {
				slog.Debug("llama runner terminated", "error", err)
				if strings.Contains(s.status.LastErrMsg, "unknown model") {
					s.status.LastErrMsg = "this model is not supported by your version of Ollama. You may need to upgrade"
				}
Michael Yang's avatar
lint  
Michael Yang committed
424
				s.done <- errors.New(s.status.LastErrMsg)
425
426
427
			} else {
				s.done <- err
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
428
429
		}()

430
431
432
433
434
435
436
437
438
439
440
		return s, nil
	}

	slog.Error("unable to load any llama server", "error", finalErr)
	return nil, finalErr
}

type ServerStatus int

const ( // iota is reset to 0
	ServerStatusReady ServerStatus = iota
441
	ServerStatusNoSlotsAvailable
442
443
444
445
446
	ServerStatusLoadingModel
	ServerStatusNotResponding
	ServerStatusError
)

Daniel Hiltgen's avatar
Daniel Hiltgen committed
447
448
449
450
func (s ServerStatus) ToString() string {
	switch s {
	case ServerStatusReady:
		return "llm server ready"
451
	case ServerStatusNoSlotsAvailable:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
452
453
454
455
456
457
458
459
460
461
		return "llm busy - no slots available"
	case ServerStatusLoadingModel:
		return "llm server loading model"
	case ServerStatusNotResponding:
		return "llm server not responding"
	default:
		return "llm server error"
	}
}

462
type ServerStatusResp struct {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
463
464
465
466
467
	Status          string  `json:"status"`
	SlotsIdle       int     `json:"slots_idle"`
	SlotsProcessing int     `json:"slots_processing"`
	Error           string  `json:"error"`
	Progress        float32 `json:"progress"`
468
469
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
470
func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
471
472
473
474
475
476
	// Fail fast if its exited
	if s.cmd.ProcessState != nil {
		msg := ""
		if s.status != nil && s.status.LastErrMsg != "" {
			msg = s.status.LastErrMsg
		}
477
478
479
480
		if s.cmd.ProcessState.ExitCode() == -1 {
			// Most likely a signal killed it, log some more details to try to help troubleshoot
			slog.Warn("llama runner process no longer running", "sys", s.cmd.ProcessState.Sys(), "string", s.cmd.ProcessState.String())
		}
481
482
483
484
485
486
487
488
489
490
491
492
		return ServerStatusError, fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d/health", s.port), nil)
	if err != nil {
		return ServerStatusError, fmt.Errorf("error creating GET request: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
Michael Yang's avatar
Michael Yang committed
493
			return ServerStatusNotResponding, errors.New("server not responding")
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
		}
		return ServerStatusError, fmt.Errorf("health resp: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return ServerStatusError, fmt.Errorf("read health request: %w", err)
	}

	var status ServerStatusResp
	if err := json.Unmarshal(body, &status); err != nil {
		return ServerStatusError, fmt.Errorf("health unmarshal encode response: %w", err)
	}

	switch status.Status {
	case "ok":
		return ServerStatusReady, nil
	case "no slot available":
513
		return ServerStatusNoSlotsAvailable, nil
514
	case "loading model":
Daniel Hiltgen's avatar
Daniel Hiltgen committed
515
		s.loadProgress = status.Progress
516
517
518
519
520
521
		return ServerStatusLoadingModel, nil
	default:
		return ServerStatusError, fmt.Errorf("server error: %+v", status)
	}
}

522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
// getServerStatusRetry will retry if ServerStatusNoSlotsAvailable is received
func (s *llmServer) getServerStatusRetry(ctx context.Context) (ServerStatus, error) {
	var retries int
	for {
		status, err := s.getServerStatus(ctx)
		if err != nil {
			return status, err
		}

		if status == ServerStatusNoSlotsAvailable {
			if retries >= 10 {
				return status, fmt.Errorf("no slots available after %d retries", retries)
			}

			time.Sleep(5 * time.Millisecond)
			retries++
			continue
		}

		return status, nil
	}
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
545
func (s *llmServer) Ping(ctx context.Context) error {
546
547
548
549
550
551
552
553
	_, err := s.getServerStatus(ctx)
	if err != nil {
		slog.Debug("server unhealthy", "error", err)
		return err
	}
	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
554
func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
555
	start := time.Now()
556
	stallDuration := envconfig.LoadTimeout()    // If no progress happens
557
	stallTimer := time.Now().Add(stallDuration) // give up if we stall
558
559
560

	slog.Info("waiting for llama runner to start responding")
	var lastStatus ServerStatus = -1
561
	fullyLoaded := false
ManniX-ITA's avatar
ManniX-ITA committed
562

563
564
	for {
		select {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
565
		case <-ctx.Done():
566
			slog.Warn("client connection closed before server finished loading, aborting load")
567
			return fmt.Errorf("timed out waiting for llama runner to start: %w", ctx.Err())
568
		case err := <-s.done:
569
			return fmt.Errorf("llama runner process has terminated: %w", err)
570
571
		default:
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
572
		if time.Now().After(stallTimer) {
ManniX-ITA's avatar
ManniX-ITA committed
573
			// timeout
574
575
576
577
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
578
			return fmt.Errorf("timed out waiting for llama runner to start - progress %0.2f - %s", s.loadProgress, msg)
ManniX-ITA's avatar
ManniX-ITA committed
579
580
581
582
583
		}
		if s.cmd.ProcessState != nil {
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
584
			}
ManniX-ITA's avatar
ManniX-ITA committed
585
586
			return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
587
588
		ctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
		defer cancel()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
589
		priorProgress := s.loadProgress
Daniel Hiltgen's avatar
Daniel Hiltgen committed
590
591
592
593
594
		status, _ := s.getServerStatus(ctx)
		if lastStatus != status && status != ServerStatusReady {
			// Only log on status changes
			slog.Info("waiting for server to become available", "status", status.ToString())
		}
ManniX-ITA's avatar
ManniX-ITA committed
595
596
		switch status {
		case ServerStatusReady:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
597
598
			s.loadDuration = time.Since(start)
			slog.Info(fmt.Sprintf("llama runner started in %0.2f seconds", s.loadDuration.Seconds()))
ManniX-ITA's avatar
ManniX-ITA committed
599
600
			return nil
		default:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
601
			lastStatus = status
Daniel Hiltgen's avatar
Daniel Hiltgen committed
602
603
604
605
			// Reset the timer as long as we're making forward progress on the load
			if priorProgress != s.loadProgress {
				slog.Debug(fmt.Sprintf("model load progress %0.2f", s.loadProgress))
				stallTimer = time.Now().Add(stallDuration)
606
607
			} else if !fullyLoaded && int(s.loadProgress*100.0) >= 100 {
				slog.Debug("model load completed, waiting for server to become available", "status", status.ToString())
608
				stallTimer = time.Now().Add(stallDuration)
609
				fullyLoaded = true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
610
			}
ManniX-ITA's avatar
ManniX-ITA committed
611
612
			time.Sleep(time.Millisecond * 250)
			continue
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
		}
	}
}

const jsonGrammar = `
root   ::= object
value  ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
635
    [^"\\\x7F\x00-\x1F] |
636
637
638
639
640
641
642
643
644
645
646
647
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

# Optional space: by convention, applied in this grammar after literal chars when allowed
ws ::= ([ \t\n] ws)?
`

const maxBufferSize = 512 * format.KiloByte

type ImageData struct {
648
649
650
	Data          []byte `json:"data"`
	ID            int    `json:"id"`
	AspectRatioID int    `json:"aspect_ratio_id"`
651
652
653
}

type completion struct {
654
655
656
657
658
	Content      string `json:"content"`
	Model        string `json:"model"`
	Prompt       string `json:"prompt"`
	Stop         bool   `json:"stop"`
	StoppedLimit bool   `json:"stopped_limit"`
659
660
661
662
663
664
665
666
667
668
669
670
671

	Timings struct {
		PredictedN  int     `json:"predicted_n"`
		PredictedMS float64 `json:"predicted_ms"`
		PromptN     int     `json:"prompt_n"`
		PromptMS    float64 `json:"prompt_ms"`
	}
}

type CompletionRequest struct {
	Prompt  string
	Format  string
	Images  []ImageData
Michael Yang's avatar
Michael Yang committed
672
	Options *api.Options
673
674
675
676
}

type CompletionResponse struct {
	Content            string
677
	DoneReason         string
678
679
680
681
682
683
684
	Done               bool
	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
685
686
func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
	if err := s.sem.Acquire(ctx, 1); err != nil {
687
688
689
690
691
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
			slog.Error("Failed to acquire semaphore", "error", err)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
692
693
694
		return err
	}
	defer s.sem.Release(1)
695

696
	// put an upper limit on num_predict to avoid the model running on forever
697
698
699
700
	if req.Options.NumPredict < 0 || req.Options.NumPredict > 10*s.options.NumCtx {
		req.Options.NumPredict = 10 * s.options.NumCtx
	}

701
702
703
704
705
706
707
708
709
	request := map[string]any{
		"prompt":            req.Prompt,
		"stream":            true,
		"n_predict":         req.Options.NumPredict,
		"n_keep":            req.Options.NumKeep,
		"main_gpu":          req.Options.MainGPU,
		"temperature":       req.Options.Temperature,
		"top_k":             req.Options.TopK,
		"top_p":             req.Options.TopP,
710
		"min_p":             req.Options.MinP,
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
		"tfs_z":             req.Options.TFSZ,
		"typical_p":         req.Options.TypicalP,
		"repeat_last_n":     req.Options.RepeatLastN,
		"repeat_penalty":    req.Options.RepeatPenalty,
		"presence_penalty":  req.Options.PresencePenalty,
		"frequency_penalty": req.Options.FrequencyPenalty,
		"mirostat":          req.Options.Mirostat,
		"mirostat_tau":      req.Options.MirostatTau,
		"mirostat_eta":      req.Options.MirostatEta,
		"penalize_nl":       req.Options.PenalizeNewline,
		"seed":              req.Options.Seed,
		"stop":              req.Options.Stop,
		"image_data":        req.Images,
		"cache_prompt":      true,
	}

	// Make sure the server is ready
728
	status, err := s.getServerStatusRetry(ctx)
729
730
731
	if err != nil {
		return err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
732
		return fmt.Errorf("unexpected server status: %s", status.ToString())
733
734
735
736
737
738
739
740
741
	}

	if req.Format == "json" {
		request["grammar"] = jsonGrammar
		if !strings.Contains(strings.ToLower(req.Prompt), "json") {
			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.")
		}
	}

742
743
744
745
	// Handling JSON marshaling with special characters unescaped.
	buffer := &bytes.Buffer{}
	enc := json.NewEncoder(buffer)
	enc.SetEscapeHTML(false)
746

747
748
749
	if err := enc.Encode(request); err != nil {
		return fmt.Errorf("failed to marshal data: %v", err)
	}
750

751
752
753
754
755
756
	endpoint := fmt.Sprintf("http://127.0.0.1:%d/completion", s.port)
	serverReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, buffer)
	if err != nil {
		return fmt.Errorf("error creating POST request: %v", err)
	}
	serverReq.Header.Set("Content-Type", "application/json")
757

758
759
760
761
762
	res, err := http.DefaultClient.Do(serverReq)
	if err != nil {
		return fmt.Errorf("POST predict: %v", err)
	}
	defer res.Body.Close()
763

764
765
	if res.StatusCode >= 400 {
		bodyBytes, err := io.ReadAll(res.Body)
766
		if err != nil {
767
			return fmt.Errorf("failed reading llm error response: %w", err)
768
		}
769
770
771
		log.Printf("llm predict error: %s", bodyBytes)
		return fmt.Errorf("%s", bodyBytes)
	}
772

773
774
775
	scanner := bufio.NewScanner(res.Body)
	buf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(buf, maxBufferSize)
776

777
778
779
	// keep track of the last token generated, this is used to abort if the model starts looping
	var lastToken string
	var tokenRepeat int
780

781
782
783
784
785
786
787
788
789
790
	for scanner.Scan() {
		select {
		case <-ctx.Done():
			// This handles the request cancellation
			return ctx.Err()
		default:
			line := scanner.Bytes()
			if len(line) == 0 {
				continue
			}
791

792
			// slog.Debug("got line", "line", string(line))
793
794
			evt, ok := bytes.CutPrefix(line, []byte("data: "))
			if !ok {
795
				evt = line
796
			}
797

798
799
			var c completion
			if err := json.Unmarshal(evt, &c); err != nil {
800
				return fmt.Errorf("error unmarshalling llm prediction response: %v", err)
801
802
803
804
805
806
807
808
			}
			switch {
			case strings.TrimSpace(c.Content) == lastToken:
				tokenRepeat++
			default:
				lastToken = strings.TrimSpace(c.Content)
				tokenRepeat = 0
			}
809

810
811
812
813
814
			// 30 picked as an arbitrary max token repeat limit, modify as needed
			if tokenRepeat > 30 {
				slog.Debug("prediction aborted, token repeat limit reached")
				return ctx.Err()
			}
815

816
817
818
819
820
			if c.Content != "" {
				fn(CompletionResponse{
					Content: c.Content,
				})
			}
821

822
			if c.Stop {
823
824
825
826
827
				doneReason := "stop"
				if c.StoppedLimit {
					doneReason = "length"
				}

828
829
				fn(CompletionResponse{
					Done:               true,
830
					DoneReason:         doneReason,
831
832
833
834
835
836
					PromptEvalCount:    c.Timings.PromptN,
					PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
					EvalCount:          c.Timings.PredictedN,
					EvalDuration:       parseDurationMs(c.Timings.PredictedMS),
				})
				return nil
837
838
			}
		}
839
	}
840

841
	if err := scanner.Err(); err != nil {
842
		if strings.Contains(err.Error(), "unexpected EOF") || strings.Contains(err.Error(), "forcibly closed") {
843
			s.Close()
844
			var msg string
845
846
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
847
848
			} else {
				msg = err.Error()
849
			}
850
			return fmt.Errorf("an error was encountered while running the model: %s", msg)
851
852
		}

853
		return fmt.Errorf("error reading llm response: %v", err)
854
855
	}

856
	return nil
857
858
}

859
860
type EmbeddingRequest struct {
	Content string `json:"content"`
861
862
}

863
864
type EmbeddingResponse struct {
	Embedding []float32 `json:"embedding"`
865
866
}

867
868
func (s *llmServer) Embedding(ctx context.Context, input string) ([]float32, error) {
	if err := s.sem.Acquire(ctx, 1); err != nil {
869
870
871
872
873
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embedding request due to client closing the connection")
		} else {
			slog.Error("Failed to acquire semaphore", "error", err)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
874
875
		return nil, err
	}
876
	defer s.sem.Release(1)
877

878
	// Make sure the server is ready
879
	status, err := s.getServerStatusRetry(ctx)
880
881
882
	if err != nil {
		return nil, err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
883
		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
884
885
	}

886
	data, err := json.Marshal(EmbeddingRequest{Content: input})
Michael Yang's avatar
Michael Yang committed
887
	if err != nil {
888
889
890
		return nil, fmt.Errorf("error marshaling embed data: %w", err)
	}

891
	r, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/embedding", s.port), bytes.NewBuffer(data))
892
893
894
	if err != nil {
		return nil, fmt.Errorf("error creating embed request: %w", err)
	}
895
	r.Header.Set("Content-Type", "application/json")
896

897
	resp, err := http.DefaultClient.Do(r)
898
899
900
901
902
903
904
905
906
907
908
	if err != nil {
		return nil, fmt.Errorf("do embedding request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading embed response: %w", err)
	}

	if resp.StatusCode >= 400 {
909
		log.Printf("llm embedding error: %s", body)
910
911
912
		return nil, fmt.Errorf("%s", body)
	}

913
	var e EmbeddingResponse
914
	if err := json.Unmarshal(body, &e); err != nil {
915
916
917
		return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
	}

918
	return e.Embedding, nil
919
920
}

Michael Yang's avatar
Michael Yang committed
921
922
923
924
925
926
927
928
type TokenizeRequest struct {
	Content string `json:"content"`
}

type TokenizeResponse struct {
	Tokens []int `json:"tokens"`
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
929
func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error) {
930
931
932
933
934
935
	s.modelLock.Lock()
	defer s.modelLock.Unlock()
	if s.model != nil {
		return s.model.Tokenize(content, false, true)
	}

Michael Yang's avatar
Michael Yang committed
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
	// Make sure the server is ready
	status, err := s.getServerStatus(ctx)
	if err != nil {
		return nil, err
	} else if status != ServerStatusReady && status != ServerStatusNoSlotsAvailable {
		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
	}

	data, err := json.Marshal(TokenizeRequest{Content: content})
	if err != nil {
		return nil, fmt.Errorf("marshaling encode data: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/tokenize", s.port), bytes.NewBuffer(data))
	if err != nil {
		return nil, fmt.Errorf("encode request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("do encode request: %w", err)
	}
	defer resp.Body.Close()
960
961
962
	if resp.StatusCode == http.StatusNotFound {
		if s.model == nil {
			slog.Debug("new runner detected, loading model for cgo tokenization")
963
964
965
966
			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
			if err != nil {
				return nil, err
			}
967
968
969
970
			s.model = m
		}
		return s.model.Tokenize(content, false, true)
	}
Michael Yang's avatar
Michael Yang committed
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read encode request: %w", err)
	}

	if resp.StatusCode >= 400 {
		log.Printf("llm encode error: %s", body)
		return nil, fmt.Errorf("%s", body)
	}

	var encoded TokenizeResponse
	if err := json.Unmarshal(body, &encoded); err != nil {
		return nil, fmt.Errorf("unmarshal encode response: %w", err)
	}

	return encoded.Tokens, nil
}

type DetokenizeRequest struct {
	Tokens []int `json:"tokens"`
}

type DetokenizeResponse struct {
	Content string `json:"content"`
996
997
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
998
func (s *llmServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
999
1000
1001
1002
1003
1004
1005
1006
1007
	s.modelLock.Lock()
	defer s.modelLock.Unlock()
	if s.model != nil {
		var resp string
		for _, token := range tokens {
			resp += s.model.TokenToPiece(token)
		}
		return resp, nil
	}
Michael Yang's avatar
Michael Yang committed
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
	// Make sure the server is ready
	status, err := s.getServerStatus(ctx)
	if err != nil {
		return "", err
	} else if status != ServerStatusReady && status != ServerStatusNoSlotsAvailable {
		return "", fmt.Errorf("unexpected server status: %s", status.ToString())
	}

	data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
	if err != nil {
		return "", fmt.Errorf("marshaling decode data: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/detokenize", s.port), bytes.NewBuffer(data))
	if err != nil {
		return "", fmt.Errorf("decode request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("do decode request: %w", err)
	}
	defer resp.Body.Close()
1032
1033
1034
	if resp.StatusCode == http.StatusNotFound {
		if s.model == nil {
			slog.Debug("new runner detected, loading model for cgo tokenization")
1035
1036
1037
1038
			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
			if err != nil {
				return "", err
			}
1039
1040
1041
1042
1043
1044
1045
1046
			s.model = m
		}
		var resp string
		for _, token := range tokens {
			resp += s.model.TokenToPiece(token)
		}
		return resp, nil
	}
Michael Yang's avatar
Michael Yang committed
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("read decode request: %w", err)
	}

	if resp.StatusCode >= 400 {
		log.Printf("llm decode error: %s", body)
		return "", fmt.Errorf("%s", body)
	}

	var decoded DetokenizeResponse
	if err := json.Unmarshal(body, &decoded); err != nil {
		return "", fmt.Errorf("unmarshal encode response: %w", err)
	}

	return decoded.Content, nil
1064
1065
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1066
func (s *llmServer) Close() error {
1067
	s.modelLock.Lock()
1068
1069
1070
1071
	if s.model != nil {
		llama.FreeModel(s.model)
		s.model = nil
	}
1072
1073
	s.modelLock.Unlock()

1074
1075
	if s.cmd != nil {
		slog.Debug("stopping llama server")
1076
1077
1078
		if err := s.cmd.Process.Kill(); err != nil {
			return err
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1079
1080
1081
1082
1083
		// if ProcessState is already populated, Wait already completed, no need to wait again
		if s.cmd.ProcessState == nil {
			slog.Debug("waiting for llama server to exit")
			<-s.done
		}
1084
1085

		slog.Debug("llama server stopped")
1086
1087
1088
1089
1090
	}

	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1091
func (s *llmServer) EstimatedVRAM() uint64 {
1092
	return s.estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1093
1094
}

1095
func (s *llmServer) EstimatedTotal() uint64 {
1096
	return s.estimate.TotalSize
1097
1098
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1099
func (s *llmServer) EstimatedVRAMByGPU(gpuID string) uint64 {
1100
1101
	for i, gpu := range s.gpus {
		if gpu.ID == gpuID {
1102
1103
1104
			if i < len(s.estimate.GPUSizes) {
				return s.estimate.GPUSizes[i]
			}
1105
1106
1107
1108
1109
		}
	}
	return 0
}

1110
1111
1112
1113
1114
1115
1116
1117
func parseDurationMs(ms float64) time.Duration {
	dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
	if err != nil {
		panic(err)
	}

	return dur
}