server.go 32.2 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"
Michael Yang's avatar
Michael Yang committed
31
	"github.com/ollama/ollama/fs/ggml"
32
	"github.com/ollama/ollama/llama"
33
34
)

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

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

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

	sem *semaphore.Weighted
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.
Michael Yang's avatar
Michael Yang committed
75
func LoadModel(model string, maxArraySize int) (*ggml.GGML, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
76
77
78
79
	if _, err := os.Stat(model); err != nil {
		return nil, err
	}

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

Michael Yang's avatar
Michael Yang committed
86
	ggml, _, err := ggml.Decode(f, maxArraySize)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
87
88
	return ggml, err
}
89

Daniel Hiltgen's avatar
Daniel Hiltgen committed
90
91
// NewLlamaServer will run a server for the given GPUs
// The gpu list must be a single family.
Michael Yang's avatar
Michael Yang committed
92
func NewLlamaServer(gpus discover.GpuInfoList, model string, f *ggml.GGML, adapters, projectors []string, opts api.Options, numParallel int) (LlamaServer, error) {
93
	systemInfo := discover.GetSystemInfo()
Michael Yang's avatar
Michael Yang committed
94
95
96
	systemTotalMemory := systemInfo.System.TotalMemory
	systemFreeMemory := systemInfo.System.FreeMemory
	systemSwapFreeMemory := systemInfo.System.FreeSwap
97
	slog.Info("system memory", "total", format.HumanBytes2(systemTotalMemory), "free", format.HumanBytes2(systemFreeMemory), "free_swap", format.HumanBytes2(systemSwapFreeMemory))
98

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

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

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

Michael Yang's avatar
Michael Yang committed
130
	slog.Info("offload", "", estimate)
131

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

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

Michael Yang's avatar
Michael Yang committed
142
	if envconfig.Debug() {
143
144
145
146
		params = append(params, "--verbose")
	}

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
174
	if fa && !f.SupportsFlashAttention() {
175
176
177
178
		slog.Warn("flash attention enabled but not supported by model")
		fa = false
	}

179
	kvct := strings.ToLower(envconfig.KvCacheType())
180
181
182
183
184
185
186

	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
Michael Yang's avatar
Michael Yang committed
187
		if kvct != "" && f.SupportsKVCacheType(kvct) {
188
189
190
			params = append(params, "--kv-cache-type", kvct)
		} else {
			slog.Warn("kv cache type not supported by model", "type", kvct)
Sam's avatar
Sam committed
191
		}
192
193
194
	} else if kvct != "" && kvct != "f16" {
		slog.Warn("quantized kv cache requested but flash attention disabled", "type", kvct)
	}
195

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

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

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

220
	// TODO - NUMA support currently doesn't work properly
221

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

224
225
226
227
	if estimate.TensorSplit != "" {
		params = append(params, "--tensor-split", estimate.TensorSplit)
	}

228
229
230
231
	if envconfig.MultiUserCache() {
		params = append(params, "--multiuser-cache")
	}

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

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

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

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

338
		s.cmd.Env = os.Environ()
339
340
		s.cmd.Stdout = os.Stdout
		s.cmd.Stderr = s.status
341
		s.cmd.SysProcAttr = LlamaServerSysProcAttr
342

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

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

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

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

427
428
429
430
431
432
433
434
		return s, nil
	}
}

type ServerStatus int

const ( // iota is reset to 0
	ServerStatusReady ServerStatus = iota
435
	ServerStatusNoSlotsAvailable
436
437
438
439
440
	ServerStatusLoadingModel
	ServerStatusNotResponding
	ServerStatusError
)

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

456
type ServerStatusResp struct {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
457
458
459
460
461
	Status          string  `json:"status"`
	SlotsIdle       int     `json:"slots_idle"`
	SlotsProcessing int     `json:"slots_processing"`
	Error           string  `json:"error"`
	Progress        float32 `json:"progress"`
462
463
}

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

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

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

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

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

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

type completion struct {
643
644
645
646
647
	Content      string `json:"content"`
	Model        string `json:"model"`
	Prompt       string `json:"prompt"`
	Stop         bool   `json:"stop"`
	StoppedLimit bool   `json:"stopped_limit"`
648
649
650
651
652
653
654
655
656
657
658

	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
659
	Format  json.RawMessage
660
	Images  []ImageData
Michael Yang's avatar
Michael Yang committed
661
	Options *api.Options
662
663
664
665
}

type CompletionResponse struct {
	Content            string
666
	DoneReason         string
667
668
669
670
671
672
673
	Done               bool
	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration
}

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

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

712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
			// 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
	}

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

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

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

753
754
755
756
757
758
	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")
759

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

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

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

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

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

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

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

812
813
814
815
816
			// 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()
			}
817

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

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

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

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

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

858
	return nil
859
860
}

861
862
type EmbeddingRequest struct {
	Content string `json:"content"`
863
864
}

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

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

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

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

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

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

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

920
	return e.Embedding, nil
921
922
}

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

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

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

	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"`
998
999
}

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

	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
1066
1067
}

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

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

		slog.Debug("llama server stopped")
1088
1089
1090
1091
1092
	}

	return nil
}

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

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

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

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

	return dur
}