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) {
92
93
	var systemTotalMemory uint64
	var systemFreeMemory uint64
94
	var systemSwapFreeMemory uint64
95

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

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

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

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

133
134
	estimate.log()

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

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

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

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

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

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

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

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

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

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

	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
194
		}
195
196
197
	} else if kvct != "" && kvct != "f16" {
		slog.Warn("quantized kv cache requested but flash attention disabled", "type", kvct)
	}
198

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

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

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
235
	libs := make(map[string]string)
236
237
	if entries, err := os.ReadDir(discover.LibOllamaPath); err == nil {
		for _, entry := range entries {
Michael Yang's avatar
Michael Yang committed
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
			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...)
254
255
			continue
		}
256

Michael Yang's avatar
Michael Yang committed
257
258
259
		// 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
260
		}
Michael Yang's avatar
Michael Yang committed
261
262
	}
	slog.Debug("compatible gpu libraries", "compatible", compatible)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
263

Michael Yang's avatar
Michael Yang committed
264
265
266
267
	// 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 {
268
269
270
271
272
273
274
275
276
		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 {
277
			slog.Debug("ResolveTCPAddr failed, using random port")
278
279
			port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
		}
280
281
282
		finalParams := []string{"runner"}
		finalParams = append(finalParams, params...)
		finalParams = append(finalParams, "--port", strconv.Itoa(port))
283

284
285
286
		var pathEnv string
		switch runtime.GOOS {
		case "windows":
287
			pathEnv = "PATH"
288
289
290
291
		case "darwin":
			pathEnv = "DYLD_LIBRARY_PATH"
		default:
			pathEnv = "LD_LIBRARY_PATH"
292
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
293

Michael Yang's avatar
Michael Yang committed
294
		var libraryPaths []string
295
		if libraryPath, ok := os.LookupEnv(pathEnv); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
296
			libraryPaths = append(libraryPaths, filepath.SplitList(libraryPath)...)
297
298
		}

Michael Yang's avatar
Michael Yang committed
299
300
301
302
303
304
305
306
		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
307
		// Note: we always put the dependency path first
Daniel Hiltgen's avatar
Daniel Hiltgen committed
308
		// since this was the exact version we compiled/linked against
309
		if gpus[0].DependencyPath != nil {
Michael Yang's avatar
Michael Yang committed
310
			slog.Debug("adding gpu dependency paths", "paths", gpus[0].DependencyPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
311
			// assume gpus from the same library have the same dependency path
312
			libraryPaths = append(gpus[0].DependencyPath, libraryPaths...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
313
314
		}

Michael Yang's avatar
Michael Yang committed
315
316
317
318
319
320
321
322
323
324
325
326
327
		// 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)
		}

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

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

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

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

		if err = s.cmd.Start(); err != nil {
Michael Yang's avatar
Michael Yang committed
403
			var msg string
404
405
406
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
Michael Yang's avatar
Michael Yang committed
407
408
409
410
411
412
413
			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:]
414
415
416
			continue
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
417
418
		// reap subprocess when it exits
		go func() {
419
420
421
			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
422
				slog.Error("llama runner terminated", "error", err)
423
424
425
				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
426
				s.done <- errors.New(s.status.LastErrMsg)
427
428
429
			} else {
				s.done <- err
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
430
431
		}()

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

type ServerStatus int

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

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

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

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

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

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

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

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

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

type completion struct {
648
649
650
651
652
	Content      string `json:"content"`
	Model        string `json:"model"`
	Prompt       string `json:"prompt"`
	Stop         bool   `json:"stop"`
	StoppedLimit bool   `json:"stopped_limit"`
653
654
655
656
657
658
659
660
661
662
663

	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
664
	Format  json.RawMessage
665
	Images  []ImageData
Michael Yang's avatar
Michael Yang committed
666
	Options *api.Options
667
668
669
670
}

type CompletionResponse struct {
	Content            string
671
	DoneReason         string
672
673
674
675
676
677
678
	Done               bool
	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
679
func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
680
681
682
683
684
685
686
687
688
	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,
689
		"min_p":             req.Options.MinP,
690
691
692
693
694
695
696
697
698
699
700
701
702
703
		"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,
	}

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

717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
			// 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
	}

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

749
750
751
752
	// Handling JSON marshaling with special characters unescaped.
	buffer := &bytes.Buffer{}
	enc := json.NewEncoder(buffer)
	enc.SetEscapeHTML(false)
753

754
755
756
	if err := enc.Encode(request); err != nil {
		return fmt.Errorf("failed to marshal data: %v", err)
	}
757

758
759
760
761
762
763
	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")
764

765
766
767
768
769
	res, err := http.DefaultClient.Do(serverReq)
	if err != nil {
		return fmt.Errorf("POST predict: %v", err)
	}
	defer res.Body.Close()
770

771
772
	if res.StatusCode >= 400 {
		bodyBytes, err := io.ReadAll(res.Body)
773
		if err != nil {
774
			return fmt.Errorf("failed reading llm error response: %w", err)
775
		}
776
777
778
		log.Printf("llm predict error: %s", bodyBytes)
		return fmt.Errorf("%s", bodyBytes)
	}
779

780
781
782
	scanner := bufio.NewScanner(res.Body)
	buf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(buf, maxBufferSize)
783

784
785
786
	// keep track of the last token generated, this is used to abort if the model starts looping
	var lastToken string
	var tokenRepeat int
787

788
789
790
791
792
793
794
795
796
797
	for scanner.Scan() {
		select {
		case <-ctx.Done():
			// This handles the request cancellation
			return ctx.Err()
		default:
			line := scanner.Bytes()
			if len(line) == 0 {
				continue
			}
798

799
			// slog.Debug("got line", "line", string(line))
800
801
			evt, ok := bytes.CutPrefix(line, []byte("data: "))
			if !ok {
802
				evt = line
803
			}
804

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

817
818
819
820
821
			// 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()
			}
822

823
824
825
826
827
			if c.Content != "" {
				fn(CompletionResponse{
					Content: c.Content,
				})
			}
828

829
			if c.Stop {
830
831
832
833
834
				doneReason := "stop"
				if c.StoppedLimit {
					doneReason = "length"
				}

835
836
				fn(CompletionResponse{
					Done:               true,
837
					DoneReason:         doneReason,
838
839
840
841
842
843
					PromptEvalCount:    c.Timings.PromptN,
					PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
					EvalCount:          c.Timings.PredictedN,
					EvalDuration:       parseDurationMs(c.Timings.PredictedMS),
				})
				return nil
844
845
			}
		}
846
	}
847

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

860
		return fmt.Errorf("error reading llm response: %v", err)
861
862
	}

863
	return nil
864
865
}

866
867
type EmbeddingRequest struct {
	Content string `json:"content"`
868
869
}

870
871
type EmbeddingResponse struct {
	Embedding []float32 `json:"embedding"`
872
873
}

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

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

893
	data, err := json.Marshal(EmbeddingRequest{Content: input})
Michael Yang's avatar
Michael Yang committed
894
	if err != nil {
895
896
897
		return nil, fmt.Errorf("error marshaling embed data: %w", err)
	}

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

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

920
	var e EmbeddingResponse
921
	if err := json.Unmarshal(body, &e); err != nil {
922
923
924
		return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
	}

925
	return e.Embedding, nil
926
927
}

Michael Yang's avatar
Michael Yang committed
928
929
930
931
932
933
934
935
type TokenizeRequest struct {
	Content string `json:"content"`
}

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

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

	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"`
1003
1004
}

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

	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
1071
1072
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1073
func (s *llmServer) Close() error {
1074
	s.modelLock.Lock()
1075
1076
1077
1078
	if s.model != nil {
		llama.FreeModel(s.model)
		s.model = nil
	}
1079
1080
	s.modelLock.Unlock()

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

		slog.Debug("llama server stopped")
1093
1094
1095
1096
1097
	}

	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1098
func (s *llmServer) EstimatedVRAM() uint64 {
1099
	return s.estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1100
1101
}

1102
func (s *llmServer) EstimatedTotal() uint64 {
1103
	return s.estimate.TotalSize
1104
1105
}

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

1117
1118
1119
1120
1121
1122
1123
1124
func parseDurationMs(ms float64) time.Duration {
	dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
	if err != nil {
		panic(err)
	}

	return dur
}