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/discover"
29
	"github.com/ollama/ollama/envconfig"
30
	"github.com/ollama/ollama/format"
31
	"github.com/ollama/ollama/llama"
32
33
)

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

// llmServer is an instance of the llama.cpp server
type llmServer struct {
49
50
51
52
53
54
	port        int
	cmd         *exec.Cmd
	done        chan error // Channel to signal when the process exits
	status      *StatusWriter
	options     api.Options
	numParallel int
55
56
57
	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
58

59
60
61
	estimate    MemoryEstimate
	totalLayers uint64
	// gpuCount     int
62
63
	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
64
	loadProgress float32
Daniel Hiltgen's avatar
Daniel Hiltgen committed
65
66

	sem *semaphore.Weighted
67
68
}

69
70
71
72
73
74
// 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
75
76
77
78
	if _, err := os.Stat(model); err != nil {
		return nil, err
	}

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

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
89
90
// NewLlamaServer will run a server for the given GPUs
// The gpu list must be a single family.
91
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
92
	var err error
93
94
	var systemTotalMemory uint64
	var systemFreeMemory uint64
95
	var systemSwapFreeMemory uint64
96

97
	systemInfo := discover.GetSystemInfo()
98
99
100
101
	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))
102

103
104
	// If the user wants zero GPU layers, reset the gpu list to be CPU/system ram info
	if opts.NumGPU == 0 {
105
		gpus = discover.GetCPUInfo()
106
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
107

Michael Yang's avatar
Michael Yang committed
108
109
	estimate := EstimateGPULayers(gpus, ggml, projectors, opts)
	if len(gpus) > 1 || gpus[0].Library != "cpu" {
Michael Yang's avatar
Michael Yang committed
110
		switch {
111
		case gpus[0].Library == "metal" && estimate.VRAMSize > systemTotalMemory:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
112
113
114
			// disable partial offloading when model is greater than total system memory as this
			// can lead to locking up the system
			opts.NumGPU = 0
115
		case gpus[0].Library != "metal" && estimate.Layers == 0:
116
			// Don't bother loading into the GPU if no layers can fit
117
			gpus = discover.GetCPUInfo()
118
119
		case opts.NumGPU < 0 && estimate.Layers > 0 && gpus[0].Library != "cpu":
			opts.NumGPU = estimate.Layers
120
121
122
		}
	}

123
124
125
	// 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" {
126
		systemMemoryRequired := estimate.TotalSize - estimate.VRAMSize
127
		available := systemFreeMemory + systemSwapFreeMemory
128
129
130
		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))
131
132
133
		}
	}

134
135
	estimate.log()

136
137
	params := []string{
		"--model", model,
Michael Yang's avatar
lint  
Michael Yang committed
138
139
		"--ctx-size", strconv.Itoa(opts.NumCtx),
		"--batch-size", strconv.Itoa(opts.NumBatch),
140
	}
Michael Yang's avatar
Michael Yang committed
141

Michael Yang's avatar
Michael Yang committed
142
	if opts.NumGPU >= 0 {
Michael Yang's avatar
lint  
Michael Yang committed
143
		params = append(params, "--n-gpu-layers", strconv.Itoa(opts.NumGPU))
144
145
	}

Michael Yang's avatar
Michael Yang committed
146
	if envconfig.Debug() {
147
148
149
150
		params = append(params, "--verbose")
	}

	if opts.MainGPU > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
151
		params = append(params, "--main-gpu", strconv.Itoa(opts.MainGPU))
152
153
154
	}

	if len(adapters) > 0 {
155
156
157
		for _, adapter := range adapters {
			params = append(params, "--lora", adapter)
		}
158
159
160
161
162
163
164
	}

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

165
	defaultThreads := systemInfo.GetOptimalThreadCount()
166
	if opts.NumThread > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
167
		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
168
169
	} else if defaultThreads > 0 {
		params = append(params, "--threads", strconv.Itoa(defaultThreads))
170
171
	}

172
173
174
175
176
	fa := envconfig.FlashAttention()
	if fa && !gpus.FlashAttentionSupported() {
		slog.Warn("flash attention enabled but not supported by gpu")
		fa = false
	}
Sam's avatar
Sam committed
177

178
179
180
181
182
	if fa && !ggml.SupportsFlashAttention() {
		slog.Warn("flash attention enabled but not supported by model")
		fa = false
	}

183
	kvct := strings.ToLower(envconfig.KvCacheType())
184
185
186
187
188
189
190
191
192
193
194

	if fa {
		slog.Info("enabling flash attention")
		params = append(params, "--flash-attn")

		// Flash Attention also supports kv cache quantization
		// Enable if the requested and kv cache type is supported by the model
		if kvct != "" && ggml.SupportsKVCacheType(kvct) {
			params = append(params, "--kv-cache-type", kvct)
		} else {
			slog.Warn("kv cache type not supported by model", "type", kvct)
Sam's avatar
Sam committed
195
		}
196
197
198
	} else if kvct != "" && kvct != "f16" {
		slog.Warn("quantized kv cache requested but flash attention disabled", "type", kvct)
	}
199

200
201
	// mmap has issues with partial offloading on metal
	for _, g := range gpus {
202
203
204
		if g.Library == "metal" &&
			uint64(opts.NumGPU) > 0 &&
			uint64(opts.NumGPU) < ggml.KV().BlockCount()+1 {
205
206
			opts.UseMMap = new(bool)
			*opts.UseMMap = false
207
		}
Sam's avatar
Sam committed
208
	}
209

210
	// Windows CUDA should not use mmap for best performance
211
	// Linux  with a model larger than free space, mmap leads to thrashing
Daniel Hiltgen's avatar
Daniel Hiltgen committed
212
	// For CPU loads we want the memory to be allocated, not FS cache
213
214
215
216
	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) {
217
218
219
220
221
222
223
		params = append(params, "--no-mmap")
	}

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

224
	// TODO - NUMA support currently doesn't work properly
225

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

228
229
230
231
	if estimate.TensorSplit != "" {
		params = append(params, "--tensor-split", estimate.TensorSplit)
	}

232
233
234
235
	if envconfig.MultiUserCache() {
		params = append(params, "--multiuser-cache")
	}

Michael Yang's avatar
Michael Yang committed
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
	// get available libraries
	if err != nil {
		return nil, fmt.Errorf("could not get libollama dir: %w", err)
	}

	entries, err := os.ReadDir(discover.LibOllamaPath)
	if err != nil {
		return nil, fmt.Errorf("could not read libollama dir: %w", err)
	}

	libs := make(map[string]string)
	for _, entry := range entries {
		if entry.IsDir() {
			libs[entry.Name()] = filepath.Join(discover.LibOllamaPath, entry.Name())
		}
	}

	lib := gpus[0].RunnerName()
	requested := envconfig.LLMLibrary()
	if libs[requested] != "" {
		slog.Info("using requested gpu library", "requested", requested)
		lib = requested
	}

	var compatible []string
	for k := range libs {
		// exact match first
		if k == lib {
			compatible = append([]string{k}, compatible...)
265
266
			continue
		}
267

Michael Yang's avatar
Michael Yang committed
268
269
270
		// then match the family (e.g. 'cuda')
		if strings.Split(k, "_")[0] == strings.Split(lib, "_")[0] {
			compatible = append(compatible, k)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
271
		}
Michael Yang's avatar
Michael Yang committed
272
273
	}
	slog.Debug("compatible gpu libraries", "compatible", compatible)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
274

Michael Yang's avatar
Michael Yang committed
275
276
277
278
	// iterate through compatible GPU libraries such as 'cuda_v12', 'cuda_v11', 'rocm', etc.
	// adding each library's respective path to the LD_LIBRARY_PATH, until finally running
	// without any LD_LIBRARY_PATH flags
	for {
279
280
281
282
283
284
285
286
287
288
289
290
		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
		}
291
292
293
		finalParams := []string{"runner"}
		finalParams = append(finalParams, params...)
		finalParams = append(finalParams, "--port", strconv.Itoa(port))
294
295
296
297
298

		pathEnv := "LD_LIBRARY_PATH"
		if runtime.GOOS == "windows" {
			pathEnv = "PATH"
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
299

Michael Yang's avatar
Michael Yang committed
300
		var libraryPaths []string
301
		if libraryPath, ok := os.LookupEnv(pathEnv); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
302
			libraryPaths = append(libraryPaths, filepath.SplitList(libraryPath)...)
303
304
		}

Michael Yang's avatar
Michael Yang committed
305
306
307
308
309
310
311
312
		if len(compatible) > 0 {
			c := compatible[0]
			if libpath, ok := libs[c]; ok {
				slog.Debug("adding gpu library", "path", libpath)
				libraryPaths = append(libraryPaths, libpath)
			}
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
313
		// Note: we always put the dependency path first
Daniel Hiltgen's avatar
Daniel Hiltgen committed
314
		// since this was the exact version we compiled/linked against
315
		if gpus[0].DependencyPath != nil {
Michael Yang's avatar
Michael Yang committed
316
			slog.Debug("adding gpu dependency paths", "paths", gpus[0].DependencyPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
317
			// assume gpus from the same library have the same dependency path
318
			libraryPaths = append(gpus[0].DependencyPath, libraryPaths...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
319
320
		}

Michael Yang's avatar
Michael Yang committed
321
322
323
324
325
326
327
328
329
330
331
332
333
		// finally, add the root library path
		libraryPaths = append(libraryPaths, discover.LibOllamaPath)

		exe, err := os.Executable()
		if err != nil {
			return nil, fmt.Errorf("unable to lookup executable path: %w", err)
		}

		exe, err = filepath.EvalSymlinks(exe)
		if err != nil {
			return nil, fmt.Errorf("unable to evaluate symlinks for executable path: %w", err)
		}

334
		// TODO - once fully switched to the Go runner, load the model here for tokenize/detokenize cgo access
Daniel Hiltgen's avatar
Daniel Hiltgen committed
335
		s := &llmServer{
336
			port:        port,
Michael Yang's avatar
Michael Yang committed
337
			cmd:         exec.Command(exe, finalParams...),
338
339
			status:      NewStatusWriter(os.Stderr),
			options:     opts,
340
			modelPath:   model,
341
			estimate:    estimate,
342
			numParallel: numParallel,
343
344
			sem:         semaphore.NewWeighted(int64(numParallel)),
			totalLayers: ggml.KV().BlockCount() + 1,
345
			gpus:        gpus,
346
			done:        make(chan error, 1),
347
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
348

349
		s.cmd.Env = os.Environ()
350
351
		s.cmd.Stdout = os.Stdout
		s.cmd.Stderr = s.status
352
		s.cmd.SysProcAttr = LlamaServerSysProcAttr
353

Daniel Hiltgen's avatar
Daniel Hiltgen committed
354
355
356
357
		envWorkarounds := [][2]string{}
		for _, gpu := range gpus {
			envWorkarounds = append(envWorkarounds, gpu.EnvWorkarounds...)
		}
Michael Yang's avatar
lint  
Michael Yang committed
358
		visibleDevicesEnv, visibleDevicesEnvVal := gpus.GetVisibleDevicesEnv()
359
360
361
362
363
364
365
366
367
368
369
370
371
		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
372
373
374
375
376
377
			} else if len(envWorkarounds) != 0 {
				for _, kv := range envWorkarounds {
					if strings.EqualFold(cmp[0], kv[0]) {
						s.cmd.Env[i] = kv[0] + "=" + kv[1]
					}
				}
378
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
379
		}
380
381
		if pathNeeded {
			s.cmd.Env = append(s.cmd.Env, pathEnv+"="+pathEnvVal)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
382
		}
383
384
		if devicesNeeded {
			s.cmd.Env = append(s.cmd.Env, visibleDevicesEnv+"="+visibleDevicesEnvVal)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
385
386
		}

387
		slog.Info("starting llama server", "cmd", s.cmd.String())
Michael Yang's avatar
Michael Yang committed
388
		if envconfig.Debug() {
389
390
391
			filteredEnv := []string{}
			for _, ev := range s.cmd.Env {
				if strings.HasPrefix(ev, "CUDA_") ||
Daniel Hiltgen's avatar
Daniel Hiltgen committed
392
					strings.HasPrefix(ev, "ROCR_") ||
393
394
					strings.HasPrefix(ev, "ROCM_") ||
					strings.HasPrefix(ev, "HIP_") ||
Daniel Hiltgen's avatar
Daniel Hiltgen committed
395
					strings.HasPrefix(ev, "GPU_") ||
396
397
398
399
400
401
402
403
404
405
					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)
		}
406
407

		if err = s.cmd.Start(); err != nil {
Michael Yang's avatar
Michael Yang committed
408
			var msg string
409
410
411
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
Michael Yang's avatar
Michael Yang committed
412
413
414
415
416
417
418
			err := fmt.Errorf("error starting runner: %v %s", err, msg)
			if len(compatible) == 0 {
				return nil, err
			}

			slog.Warn("unable to start runner with compatible gpu", "error", err, "compatible", compatible)
			compatible = compatible[1:]
419
420
421
			continue
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
422
423
		// reap subprocess when it exits
		go func() {
424
425
426
			err := s.cmd.Wait()
			// Favor a more detailed message over the process exit status
			if err != nil && s.status != nil && s.status.LastErrMsg != "" {
Michael Yang's avatar
Michael Yang committed
427
				slog.Error("llama runner terminated", "error", err)
428
429
430
				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
431
				s.done <- errors.New(s.status.LastErrMsg)
432
433
434
			} else {
				s.done <- err
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
435
436
		}()

437
438
439
440
441
442
443
444
		return s, nil
	}
}

type ServerStatus int

const ( // iota is reset to 0
	ServerStatusReady ServerStatus = iota
445
	ServerStatusNoSlotsAvailable
446
447
448
449
450
	ServerStatusLoadingModel
	ServerStatusNotResponding
	ServerStatusError
)

Daniel Hiltgen's avatar
Daniel Hiltgen committed
451
452
453
454
func (s ServerStatus) ToString() string {
	switch s {
	case ServerStatusReady:
		return "llm server ready"
455
	case ServerStatusNoSlotsAvailable:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
456
457
458
459
460
461
462
463
464
465
		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"
	}
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
474
func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
475
476
477
478
479
480
	// Fail fast if its exited
	if s.cmd.ProcessState != nil {
		msg := ""
		if s.status != nil && s.status.LastErrMsg != "" {
			msg = s.status.LastErrMsg
		}
481
482
483
484
		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())
		}
485
486
487
488
489
490
491
492
493
494
495
496
		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
497
			return ServerStatusNotResponding, errors.New("server not responding")
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
		}
		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":
517
		return ServerStatusNoSlotsAvailable, nil
518
	case "loading model":
Daniel Hiltgen's avatar
Daniel Hiltgen committed
519
		s.loadProgress = status.Progress
520
521
522
523
524
525
		return ServerStatusLoadingModel, nil
	default:
		return ServerStatusError, fmt.Errorf("server error: %+v", status)
	}
}

526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
// 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
549
func (s *llmServer) Ping(ctx context.Context) error {
550
551
552
553
554
555
556
557
	_, err := s.getServerStatus(ctx)
	if err != nil {
		slog.Debug("server unhealthy", "error", err)
		return err
	}
	return nil
}

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

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

567
568
	for {
		select {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
569
		case <-ctx.Done():
570
			slog.Warn("client connection closed before server finished loading, aborting load")
571
			return fmt.Errorf("timed out waiting for llama runner to start: %w", ctx.Err())
572
		case err := <-s.done:
573
			return fmt.Errorf("llama runner process has terminated: %w", err)
574
575
		default:
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
576
		if time.Now().After(stallTimer) {
ManniX-ITA's avatar
ManniX-ITA committed
577
			// timeout
578
579
580
581
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
582
			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
583
584
585
586
587
		}
		if s.cmd.ProcessState != nil {
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
588
			}
ManniX-ITA's avatar
ManniX-ITA committed
589
590
			return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
591
592
		ctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
		defer cancel()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
593
		priorProgress := s.loadProgress
Daniel Hiltgen's avatar
Daniel Hiltgen committed
594
595
596
597
598
		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
599
600
		switch status {
		case ServerStatusReady:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
601
602
			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
603
604
			return nil
		default:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
605
			lastStatus = status
Daniel Hiltgen's avatar
Daniel Hiltgen committed
606
607
608
609
			// 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)
610
611
			} else if !fullyLoaded && int(s.loadProgress*100.0) >= 100 {
				slog.Debug("model load completed, waiting for server to become available", "status", status.ToString())
612
				stallTimer = time.Now().Add(stallDuration)
613
				fullyLoaded = true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
614
			}
ManniX-ITA's avatar
ManniX-ITA committed
615
616
			time.Sleep(time.Millisecond * 250)
			continue
617
618
619
620
		}
	}
}

621
var grammarJSON = `
622
623
624
625
626
627
628
629
630
631
632
633
634
635
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 ::=
  "\"" (
636
    [^"\\\x7F\x00-\x1F] |
637
638
639
640
641
642
643
644
645
646
    "\\" (["\\/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 {
647
648
649
	Data          []byte `json:"data"`
	ID            int    `json:"id"`
	AspectRatioID int    `json:"aspect_ratio_id"`
650
651
652
}

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

	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
669
	Format  json.RawMessage
670
	Images  []ImageData
Michael Yang's avatar
Michael Yang committed
671
	Options *api.Options
672
673
674
675
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
684
func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
685
686
687
688
689
690
691
692
693
	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,
694
		"min_p":             req.Options.MinP,
695
696
697
698
699
700
701
702
703
704
705
706
707
708
		"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,
		"seed":              req.Options.Seed,
		"stop":              req.Options.Stop,
		"image_data":        req.Images,
		"cache_prompt":      true,
	}

709
	if len(req.Format) > 0 {
710
711
712
713
714
715
		switch string(req.Format) {
		case `null`, `""`:
			// Field was set, but "missing" a value. We accept
			// these as "not set".
			break
		case `"json"`:
716
			request["grammar"] = grammarJSON
717
718
719
720
721
		default:
			if req.Format[0] != '{' {
				return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format)
			}

722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
			// User provided a JSON schema
			g := llama.SchemaToGrammar(req.Format)
			if g == nil {
				return fmt.Errorf("invalid JSON schema in format")
			}
			request["grammar"] = string(g)
		}
	}

	if err := s.sem.Acquire(ctx, 1); err != nil {
		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)
		}
		return err
	}
	defer s.sem.Release(1)

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

746
	// Make sure the server is ready
747
	status, err := s.getServerStatusRetry(ctx)
748
749
750
	if err != nil {
		return err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
751
		return fmt.Errorf("unexpected server status: %s", status.ToString())
752
753
	}

754
755
756
757
	// Handling JSON marshaling with special characters unescaped.
	buffer := &bytes.Buffer{}
	enc := json.NewEncoder(buffer)
	enc.SetEscapeHTML(false)
758

759
760
761
	if err := enc.Encode(request); err != nil {
		return fmt.Errorf("failed to marshal data: %v", err)
	}
762

763
764
765
766
767
768
	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")
769

770
771
772
773
774
	res, err := http.DefaultClient.Do(serverReq)
	if err != nil {
		return fmt.Errorf("POST predict: %v", err)
	}
	defer res.Body.Close()
775

776
777
	if res.StatusCode >= 400 {
		bodyBytes, err := io.ReadAll(res.Body)
778
		if err != nil {
779
			return fmt.Errorf("failed reading llm error response: %w", err)
780
		}
781
782
783
		log.Printf("llm predict error: %s", bodyBytes)
		return fmt.Errorf("%s", bodyBytes)
	}
784

785
786
787
	scanner := bufio.NewScanner(res.Body)
	buf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(buf, maxBufferSize)
788

789
790
791
	// keep track of the last token generated, this is used to abort if the model starts looping
	var lastToken string
	var tokenRepeat int
792

793
794
795
796
797
798
799
800
801
802
	for scanner.Scan() {
		select {
		case <-ctx.Done():
			// This handles the request cancellation
			return ctx.Err()
		default:
			line := scanner.Bytes()
			if len(line) == 0 {
				continue
			}
803

804
			// slog.Debug("got line", "line", string(line))
805
806
			evt, ok := bytes.CutPrefix(line, []byte("data: "))
			if !ok {
807
				evt = line
808
			}
809

810
811
			var c completion
			if err := json.Unmarshal(evt, &c); err != nil {
812
				return fmt.Errorf("error unmarshalling llm prediction response: %v", err)
813
814
815
816
817
818
819
820
			}
			switch {
			case strings.TrimSpace(c.Content) == lastToken:
				tokenRepeat++
			default:
				lastToken = strings.TrimSpace(c.Content)
				tokenRepeat = 0
			}
821

822
823
824
825
826
			// 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()
			}
827

828
829
830
831
832
			if c.Content != "" {
				fn(CompletionResponse{
					Content: c.Content,
				})
			}
833

834
			if c.Stop {
835
836
837
838
839
				doneReason := "stop"
				if c.StoppedLimit {
					doneReason = "length"
				}

840
841
				fn(CompletionResponse{
					Done:               true,
842
					DoneReason:         doneReason,
843
844
845
846
847
848
					PromptEvalCount:    c.Timings.PromptN,
					PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
					EvalCount:          c.Timings.PredictedN,
					EvalDuration:       parseDurationMs(c.Timings.PredictedMS),
				})
				return nil
849
850
			}
		}
851
	}
852

853
	if err := scanner.Err(); err != nil {
854
		if strings.Contains(err.Error(), "unexpected EOF") || strings.Contains(err.Error(), "forcibly closed") {
855
			s.Close()
856
			var msg string
857
858
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
859
860
			} else {
				msg = err.Error()
861
			}
862
			return fmt.Errorf("an error was encountered while running the model: %s", msg)
863
864
		}

865
		return fmt.Errorf("error reading llm response: %v", err)
866
867
	}

868
	return nil
869
870
}

871
872
type EmbeddingRequest struct {
	Content string `json:"content"`
873
874
}

875
876
type EmbeddingResponse struct {
	Embedding []float32 `json:"embedding"`
877
878
}

879
880
func (s *llmServer) Embedding(ctx context.Context, input string) ([]float32, error) {
	if err := s.sem.Acquire(ctx, 1); err != nil {
881
882
883
884
885
		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
886
887
		return nil, err
	}
888
	defer s.sem.Release(1)
889

890
	// Make sure the server is ready
891
	status, err := s.getServerStatusRetry(ctx)
892
893
894
	if err != nil {
		return nil, err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
895
		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
896
897
	}

898
	data, err := json.Marshal(EmbeddingRequest{Content: input})
Michael Yang's avatar
Michael Yang committed
899
	if err != nil {
900
901
902
		return nil, fmt.Errorf("error marshaling embed data: %w", err)
	}

903
	r, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/embedding", s.port), bytes.NewBuffer(data))
904
905
906
	if err != nil {
		return nil, fmt.Errorf("error creating embed request: %w", err)
	}
907
	r.Header.Set("Content-Type", "application/json")
908

909
	resp, err := http.DefaultClient.Do(r)
910
911
912
913
914
915
916
917
918
919
920
	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 {
921
		log.Printf("llm embedding error: %s", body)
922
923
924
		return nil, fmt.Errorf("%s", body)
	}

925
	var e EmbeddingResponse
926
	if err := json.Unmarshal(body, &e); err != nil {
927
928
929
		return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
	}

930
	return e.Embedding, nil
931
932
}

Michael Yang's avatar
Michael Yang committed
933
934
935
936
937
938
939
940
type TokenizeRequest struct {
	Content string `json:"content"`
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
941
func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error) {
942
943
944
945
946
947
	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
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
	// 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()
972
973
974
	if resp.StatusCode == http.StatusNotFound {
		if s.model == nil {
			slog.Debug("new runner detected, loading model for cgo tokenization")
975
976
977
978
			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
			if err != nil {
				return nil, err
			}
979
980
981
982
			s.model = m
		}
		return s.model.Tokenize(content, false, true)
	}
Michael Yang's avatar
Michael Yang committed
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007

	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"`
1008
1009
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1010
func (s *llmServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
1011
1012
1013
1014
1015
1016
1017
1018
1019
	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
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
	// 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()
1044
1045
1046
	if resp.StatusCode == http.StatusNotFound {
		if s.model == nil {
			slog.Debug("new runner detected, loading model for cgo tokenization")
1047
1048
1049
1050
			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
			if err != nil {
				return "", err
			}
1051
1052
1053
1054
1055
1056
1057
1058
			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
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075

	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
1076
1077
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1078
func (s *llmServer) Close() error {
1079
	s.modelLock.Lock()
1080
1081
1082
1083
	if s.model != nil {
		llama.FreeModel(s.model)
		s.model = nil
	}
1084
1085
	s.modelLock.Unlock()

1086
1087
	if s.cmd != nil {
		slog.Debug("stopping llama server")
1088
1089
1090
		if err := s.cmd.Process.Kill(); err != nil {
			return err
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1091
1092
1093
1094
1095
		// 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
		}
1096
1097

		slog.Debug("llama server stopped")
1098
1099
1100
1101
1102
	}

	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1103
func (s *llmServer) EstimatedVRAM() uint64 {
1104
	return s.estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1105
1106
}

1107
func (s *llmServer) EstimatedTotal() uint64 {
1108
	return s.estimate.TotalSize
1109
1110
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1111
func (s *llmServer) EstimatedVRAMByGPU(gpuID string) uint64 {
1112
1113
	for i, gpu := range s.gpus {
		if gpu.ID == gpuID {
1114
1115
1116
			if i < len(s.estimate.GPUSizes) {
				return s.estimate.GPUSizes[i]
			}
1117
1118
1119
1120
1121
		}
	}
	return 0
}

1122
1123
1124
1125
1126
1127
1128
1129
func parseDurationMs(ms float64) time.Duration {
	dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
	if err != nil {
		panic(err)
	}

	return dur
}