server.go 31.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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"
	"time"

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

26
	"github.com/ollama/ollama/api"
27
	"github.com/ollama/ollama/envconfig"
28
29
30
31
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/gpu"
)

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

// llmServer is an instance of the llama.cpp server
type llmServer struct {
47
48
49
50
51
52
	port        int
	cmd         *exec.Cmd
	done        chan error // Channel to signal when the process exits
	status      *StatusWriter
	options     api.Options
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
53

54
55
56
57
58
	estimate    MemoryEstimate
	totalLayers uint64
	// gpuCount     int
	gpus         gpu.GpuInfoList // Recorded just before the model loaded, free space will be incorrect
	loadDuration time.Duration   // Record how long it took the model to load
59
	loadProgress float32
Daniel Hiltgen's avatar
Daniel Hiltgen committed
60
61

	sem *semaphore.Weighted
62
63
}

64
65
66
67
68
69
// 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
70
71
72
73
	if _, err := os.Stat(model); err != nil {
		return nil, err
	}

74
75
76
77
78
79
	f, err := os.Open(model)
	if err != nil {
		return nil, err
	}
	defer f.Close()

80
	ggml, _, err := DecodeGGML(f, maxArraySize)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
81
82
	return ggml, err
}
83

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

	systemMemInfo, err := gpu.GetCPUMem()
	if err != nil {
		slog.Error("failed to lookup system memory", "error", err)
	} else {
		systemTotalMemory = systemMemInfo.TotalMemory
		systemFreeMemory = systemMemInfo.FreeMemory
100
101
		systemSwapFreeMemory = systemMemInfo.FreeSwap
		slog.Debug("system memory", "total", format.HumanBytes2(systemTotalMemory), "free", format.HumanBytes2(systemFreeMemory), "free_swap", format.HumanBytes2(systemSwapFreeMemory))
102
	}
103

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

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

128
129
130
	// On linux, over-allocating CPU memory will almost always result in an error
	if runtime.GOOS == "linux" {
		systemMemoryRequired := estimate.TotalSize - estimate.VRAMSize
131
		available := systemFreeMemory + systemSwapFreeMemory
132
133
134
		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))
135
136
137
		}
	}

138
139
	estimate.log()

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

	if len(adapters) > 1 {
		return nil, errors.New("ollama supports only one lora adapter, but multiple were provided")
	}

147
148
149
150
151
152
153
154
155
156
157
158
159
160
	availableServers := getAvailableServers()
	if len(availableServers) == 0 {
		if runtime.GOOS != "windows" {
			slog.Warn("llama server binary disappeared, reinitializing payloads")
			err = Init()
			if err != nil {
				slog.Warn("failed to reinitialize payloads", "error", err)
				return nil, err
			}
			availableServers = getAvailableServers()
		} else {
			return nil, finalErr
		}
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
161
162
163
164
165
166
	var servers []string
	if cpuRunner != "" {
		servers = []string{cpuRunner}
	} else {
		servers = serversForGpu(gpus[0]) // All GPUs in the list are matching Library and Variant
	}
Michael Yang's avatar
string  
Michael Yang committed
167
	demandLib := envconfig.LLMLibrary()
168
169
170
171
172
173
174
	if demandLib != "" {
		serverPath := availableServers[demandLib]
		if serverPath == "" {
			slog.Info(fmt.Sprintf("Invalid OLLAMA_LLM_LIBRARY %s - not found", demandLib))
		} else {
			slog.Info("user override", "OLLAMA_LLM_LIBRARY", demandLib, "path", serverPath)
			servers = []string{demandLib}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
175
176
177
178
			if strings.HasPrefix(demandLib, "cpu") {
				// Omit the GPU flag to silence the warning
				opts.NumGPU = -1
			}
179
180
181
182
		}
	}

	if len(servers) == 0 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
183
		return nil, fmt.Errorf("no servers found for %v", gpus)
184
185
186
187
	}

	params := []string{
		"--model", model,
Michael Yang's avatar
lint  
Michael Yang committed
188
189
		"--ctx-size", strconv.Itoa(opts.NumCtx),
		"--batch-size", strconv.Itoa(opts.NumBatch),
190
191
		"--embedding",
	}
Michael Yang's avatar
Michael Yang committed
192
193

	params = append(params, "--log-disable")
194

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

Michael Yang's avatar
Michael Yang committed
199
	if envconfig.Debug() {
200
201
202
203
		params = append(params, "--verbose")
	}

	if opts.MainGPU > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
204
		params = append(params, "--main-gpu", strconv.Itoa(opts.MainGPU))
205
206
207
208
209
210
211
212
213
214
215
216
217
	}

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

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

	if opts.NumThread > 0 {
Michael Yang's avatar
lint  
Michael Yang committed
218
		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
219
220
221
222
223
224
	}

	if !opts.F16KV {
		params = append(params, "--memory-f32")
	}

Michael Yang's avatar
bool  
Michael Yang committed
225
	flashAttnEnabled := envconfig.FlashAttention()
Sam's avatar
Sam committed
226
227

	for _, g := range gpus {
228
		// only cuda (compute capability 7+) and metal support flash attention
Sam's avatar
Sam committed
229
		if g.Library != "metal" && (g.Library != "cuda" || g.DriverMajor < 7) {
230
			flashAttnEnabled = false
Sam's avatar
Sam committed
231
		}
232
233
234
235
236

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

242
	if flashAttnEnabled {
Sam's avatar
Sam committed
243
244
245
		params = append(params, "--flash-attn")
	}

246
	// Windows CUDA should not use mmap for best performance
247
	// Linux  with a model larger than free space, mmap leads to thrashing
Daniel Hiltgen's avatar
Daniel Hiltgen committed
248
	// For CPU loads we want the memory to be allocated, not FS cache
249
250
251
252
	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) {
253
254
255
256
257
258
259
		params = append(params, "--no-mmap")
	}

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

260
261
262
263
264
265
266
267
	if gpu.IsNUMA() {
		numaMode := "distribute"
		if runtime.GOOS == "linux" {
			if _, err := exec.LookPath("numactl"); err == nil {
				numaMode = "numactl"
			}
		}
		params = append(params, "--numa", numaMode)
268
269
	}

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

272
273
274
275
	if estimate.TensorSplit != "" {
		params = append(params, "--tensor-split", estimate.TensorSplit)
	}

Michael Yang's avatar
lint  
Michael Yang committed
276
	for i := range len(servers) {
277
		dir := availableServers[servers[i]]
278
279
280
		if dir == "" {
			// Shouldn't happen
			finalErr = fmt.Errorf("[%d] server %s not listed in available servers %v", i, servers[i], availableServers)
Michael Yang's avatar
Michael Yang committed
281
			slog.Error("server list inconsistent", "error", finalErr)
282
283
			continue
		}
284

Daniel Hiltgen's avatar
Daniel Hiltgen committed
285
		if strings.HasPrefix(servers[i], "cpu") {
286
			gpus = gpu.GetCPUInfo()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
287
288
		}

289
		// Find an availableServers  port, retry on each iteration in case the failure was a port conflict race
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
		port := 0
		if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
			var l *net.TCPListener
			if l, err = net.ListenTCP("tcp", a); err == nil {
				port = l.Addr().(*net.TCPAddr).Port
				l.Close()
			}
		}
		if port == 0 {
			slog.Debug("ResolveTCPAddr failed ", "error", err)
			port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
		}
		finalParams := append(params, "--port", strconv.Itoa(port))

		pathEnv := "LD_LIBRARY_PATH"
		if runtime.GOOS == "windows" {
			pathEnv = "PATH"
		}
308
309
		// prepend the server directory to LD_LIBRARY_PATH/PATH and the parent dir for common dependencies
		libraryPaths := []string{dir, filepath.Dir(dir)}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
310

311
312
313
		if libraryPath, ok := os.LookupEnv(pathEnv); ok {
			// Append our runner directory to the path
			// This will favor system libraries over our bundled library dependencies
Daniel Hiltgen's avatar
Daniel Hiltgen committed
314
			libraryPaths = append(libraryPaths, filepath.SplitList(libraryPath)...)
315
316
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
317
318
319
320
321
322
323
324
		// Note: we always put the dependency path first
		// since this was the exact version we verified for AMD GPUs
		// and we favor what the user had in their path
		if gpus[0].DependencyPath != "" {
			// TODO refine for multi-gpu support
			libraryPaths = append([]string{gpus[0].DependencyPath}, libraryPaths...)
		}

325
326
		server := filepath.Join(dir, "ollama_llama_server")
		if runtime.GOOS == "windows" {
Michael Yang's avatar
Michael Yang committed
327
			server += ".exe"
328
329
		}

330
331
332
333
334
335
336
337
338
339
340
		// Detect tmp cleaners wiping out the file
		_, err := os.Stat(server)
		if errors.Is(err, os.ErrNotExist) {
			slog.Warn("llama server disappeared, reinitializing payloads", "path", server, "error", err)
			err = Init()
			if err != nil {
				slog.Warn("failed to reinitialize payloads", "error", err)
				return nil, err
			}
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
341
		s := &llmServer{
342
343
344
345
346
			port:        port,
			cmd:         exec.Command(server, finalParams...),
			status:      NewStatusWriter(os.Stderr),
			options:     opts,
			estimate:    estimate,
347
			numParallel: numParallel,
348
349
			sem:         semaphore.NewWeighted(int64(numParallel)),
			totalLayers: ggml.KV().BlockCount() + 1,
350
			gpus:        gpus,
351
			done:        make(chan error, 1),
352
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
353

354
		s.cmd.Env = os.Environ()
355
356
		s.cmd.Stdout = os.Stdout
		s.cmd.Stderr = s.status
357
		s.cmd.SysProcAttr = LlamaServerSysProcAttr
358

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

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

		if err = s.cmd.Start(); err != nil {
413
414
415
416
417
			// Detect permission denied and augment them essage about noexec
			if errors.Is(err, os.ErrPermission) {
				finalErr = fmt.Errorf("unable to start server %w.  %s may have noexec set.  Set OLLAMA_TMPDIR for server to a writable executable directory", err, dir)
				continue
			}
418
419
420
421
422
423
424
425
426
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
			err = fmt.Errorf("error starting the external llama server: %v %s", err, msg)
			finalErr = err
			continue
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
427
428
		// reap subprocess when it exits
		go func() {
429
430
431
432
433
434
435
			err := s.cmd.Wait()
			// Favor a more detailed message over the process exit status
			if err != nil && s.status != nil && s.status.LastErrMsg != "" {
				slog.Debug("llama runner terminated", "error", err)
				if strings.Contains(s.status.LastErrMsg, "unknown model") {
					s.status.LastErrMsg = "this model is not supported by your version of Ollama. You may need to upgrade"
				}
Michael Yang's avatar
lint  
Michael Yang committed
436
				s.done <- errors.New(s.status.LastErrMsg)
437
438
439
			} else {
				s.done <- err
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
440
441
		}()

442
443
444
445
446
447
448
		return s, nil
	}

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

Michael Yang's avatar
Michael Yang committed
449
func projectorMemoryRequirements(filename string) uint64 {
450
451
452
453
454
455
	file, err := os.Open(filename)
	if err != nil {
		return 0
	}
	defer file.Close()

456
	ggml, _, err := DecodeGGML(file, 0)
457
458
459
460
	if err != nil {
		return 0
	}

Michael Yang's avatar
Michael Yang committed
461
462
463
	var mem uint64
	for _, layer := range ggml.Tensors().Layers() {
		mem += layer.size()
464
465
	}

Michael Yang's avatar
Michael Yang committed
466
	return mem
467
468
469
470
471
472
}

type ServerStatus int

const ( // iota is reset to 0
	ServerStatusReady ServerStatus = iota
473
	ServerStatusNoSlotsAvailable
474
475
476
477
478
	ServerStatusLoadingModel
	ServerStatusNotResponding
	ServerStatusError
)

Daniel Hiltgen's avatar
Daniel Hiltgen committed
479
480
481
482
func (s ServerStatus) ToString() string {
	switch s {
	case ServerStatusReady:
		return "llm server ready"
483
	case ServerStatusNoSlotsAvailable:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
484
485
486
487
488
489
490
491
492
493
		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"
	}
}

494
type ServerStatusResp struct {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
495
496
497
498
499
	Status          string  `json:"status"`
	SlotsIdle       int     `json:"slots_idle"`
	SlotsProcessing int     `json:"slots_processing"`
	Error           string  `json:"error"`
	Progress        float32 `json:"progress"`
500
501
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
502
func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
503
504
505
506
507
508
	// Fail fast if its exited
	if s.cmd.ProcessState != nil {
		msg := ""
		if s.status != nil && s.status.LastErrMsg != "" {
			msg = s.status.LastErrMsg
		}
509
510
511
512
		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())
		}
513
514
515
516
517
518
519
520
521
522
523
524
		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
525
			return ServerStatusNotResponding, errors.New("server not responding")
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
		}
		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":
545
		return ServerStatusNoSlotsAvailable, nil
546
	case "loading model":
Daniel Hiltgen's avatar
Daniel Hiltgen committed
547
		s.loadProgress = status.Progress
548
549
550
551
552
553
		return ServerStatusLoadingModel, nil
	default:
		return ServerStatusError, fmt.Errorf("server error: %+v", status)
	}
}

554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
// 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
577
func (s *llmServer) Ping(ctx context.Context) error {
578
579
580
581
582
583
584
585
	_, err := s.getServerStatus(ctx)
	if err != nil {
		slog.Debug("server unhealthy", "error", err)
		return err
	}
	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
586
func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
587
	start := time.Now()
588
589
590
	stallDuration := 5 * time.Minute            // If no progress happens
	finalLoadDuration := 5 * time.Minute        // After we hit 100%, give the runner more time to come online
	stallTimer := time.Now().Add(stallDuration) // give up if we stall
591
592
593

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

596
597
	for {
		select {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
598
		case <-ctx.Done():
599
			slog.Warn("client connection closed before server finished loading, aborting load")
600
			return fmt.Errorf("timed out waiting for llama runner to start: %w", ctx.Err())
601
		case err := <-s.done:
602
			return fmt.Errorf("llama runner process has terminated: %w", err)
603
604
		default:
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
605
		if time.Now().After(stallTimer) {
ManniX-ITA's avatar
ManniX-ITA committed
606
			// timeout
607
608
609
610
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
611
			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
612
613
614
615
616
		}
		if s.cmd.ProcessState != nil {
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
617
			}
ManniX-ITA's avatar
ManniX-ITA committed
618
619
			return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
620
621
		ctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
		defer cancel()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
622
		priorProgress := s.loadProgress
Daniel Hiltgen's avatar
Daniel Hiltgen committed
623
624
625
626
627
		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
628
629
		switch status {
		case ServerStatusReady:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
630
631
			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
632
633
			return nil
		default:
Daniel Hiltgen's avatar
Daniel Hiltgen committed
634
			lastStatus = status
Daniel Hiltgen's avatar
Daniel Hiltgen committed
635
636
637
638
			// 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)
639
640
641
642
			} else if !fullyLoaded && int(s.loadProgress*100.0) >= 100 {
				slog.Debug("model load completed, waiting for server to become available", "status", status.ToString())
				stallTimer = time.Now().Add(finalLoadDuration)
				fullyLoaded = true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
643
			}
ManniX-ITA's avatar
ManniX-ITA committed
644
645
			time.Sleep(time.Millisecond * 250)
			continue
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
		}
	}
}

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

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

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

string ::=
  "\"" (
668
    [^"\\\x7F\x00-\x1F] |
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
    "\\" (["\\/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 {
	Data []byte `json:"data"`
	ID   int    `json:"id"`
}

type completion struct {
686
687
688
689
690
	Content      string `json:"content"`
	Model        string `json:"model"`
	Prompt       string `json:"prompt"`
	Stop         bool   `json:"stop"`
	StoppedLimit bool   `json:"stopped_limit"`
691
692
693
694
695
696
697
698
699
700
701
702
703

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

type CompletionRequest struct {
	Prompt  string
	Format  string
	Images  []ImageData
Michael Yang's avatar
Michael Yang committed
704
	Options *api.Options
705
706
707
708
}

type CompletionResponse struct {
	Content            string
709
	DoneReason         string
710
711
712
713
714
715
716
	Done               bool
	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
717
718
719
720
721
722
func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
	if err := s.sem.Acquire(ctx, 1); err != nil {
		slog.Error("Failed to acquire semaphore", "error", err)
		return err
	}
	defer s.sem.Release(1)
723

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

729
730
731
732
733
734
735
736
737
	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,
738
		"min_p":             req.Options.MinP,
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
		"tfs_z":             req.Options.TFSZ,
		"typical_p":         req.Options.TypicalP,
		"repeat_last_n":     req.Options.RepeatLastN,
		"repeat_penalty":    req.Options.RepeatPenalty,
		"presence_penalty":  req.Options.PresencePenalty,
		"frequency_penalty": req.Options.FrequencyPenalty,
		"mirostat":          req.Options.Mirostat,
		"mirostat_tau":      req.Options.MirostatTau,
		"mirostat_eta":      req.Options.MirostatEta,
		"penalize_nl":       req.Options.PenalizeNewline,
		"seed":              req.Options.Seed,
		"stop":              req.Options.Stop,
		"image_data":        req.Images,
		"cache_prompt":      true,
	}

	// Make sure the server is ready
756
	status, err := s.getServerStatusRetry(ctx)
757
758
759
	if err != nil {
		return err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
760
		return fmt.Errorf("unexpected server status: %s", status.ToString())
761
762
763
764
765
766
767
768
769
	}

	if req.Format == "json" {
		request["grammar"] = jsonGrammar
		if !strings.Contains(strings.ToLower(req.Prompt), "json") {
			slog.Warn("Prompt does not specify that the LLM should response in JSON, but JSON format is expected. For best results specify that JSON is expected in the system prompt.")
		}
	}

770
771
772
773
	// Handling JSON marshaling with special characters unescaped.
	buffer := &bytes.Buffer{}
	enc := json.NewEncoder(buffer)
	enc.SetEscapeHTML(false)
774

775
776
777
	if err := enc.Encode(request); err != nil {
		return fmt.Errorf("failed to marshal data: %v", err)
	}
778

779
780
781
782
783
784
	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")
785

786
787
788
789
790
	res, err := http.DefaultClient.Do(serverReq)
	if err != nil {
		return fmt.Errorf("POST predict: %v", err)
	}
	defer res.Body.Close()
791

792
793
	if res.StatusCode >= 400 {
		bodyBytes, err := io.ReadAll(res.Body)
794
		if err != nil {
795
			return fmt.Errorf("failed reading llm error response: %w", err)
796
		}
797
798
799
		log.Printf("llm predict error: %s", bodyBytes)
		return fmt.Errorf("%s", bodyBytes)
	}
800

801
802
803
	scanner := bufio.NewScanner(res.Body)
	buf := make([]byte, 0, maxBufferSize)
	scanner.Buffer(buf, maxBufferSize)
804

805
806
807
	// keep track of the last token generated, this is used to abort if the model starts looping
	var lastToken string
	var tokenRepeat int
808

809
810
811
812
813
814
815
816
817
818
	for scanner.Scan() {
		select {
		case <-ctx.Done():
			// This handles the request cancellation
			return ctx.Err()
		default:
			line := scanner.Bytes()
			if len(line) == 0 {
				continue
			}
819

820
821
822
823
			evt, ok := bytes.CutPrefix(line, []byte("data: "))
			if !ok {
				return fmt.Errorf("error parsing llm response stream: %s", line)
			}
824

825
826
			var c completion
			if err := json.Unmarshal(evt, &c); err != nil {
827
				return fmt.Errorf("error unmarshalling llm prediction response: %v", err)
828
			}
829

830
831
832
833
834
835
836
			switch {
			case strings.TrimSpace(c.Content) == lastToken:
				tokenRepeat++
			default:
				lastToken = strings.TrimSpace(c.Content)
				tokenRepeat = 0
			}
837

838
839
840
841
842
			// 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()
			}
843

844
845
846
847
848
			if c.Content != "" {
				fn(CompletionResponse{
					Content: c.Content,
				})
			}
849

850
			if c.Stop {
851
852
853
854
855
				doneReason := "stop"
				if c.StoppedLimit {
					doneReason = "length"
				}

856
857
				fn(CompletionResponse{
					Done:               true,
858
					DoneReason:         doneReason,
859
860
861
862
863
864
					PromptEvalCount:    c.Timings.PromptN,
					PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
					EvalCount:          c.Timings.PredictedN,
					EvalDuration:       parseDurationMs(c.Timings.PredictedMS),
				})
				return nil
865
866
			}
		}
867
	}
868

869
870
871
872
873
874
	if err := scanner.Err(); err != nil {
		if strings.Contains(err.Error(), "unexpected EOF") {
			s.Close()
			msg := ""
			if s.status != nil && s.status.LastErrMsg != "" {
				msg = s.status.LastErrMsg
875
			}
876
			return fmt.Errorf("an unknown error was encountered while running the model %s", msg)
877
878
		}

879
		return fmt.Errorf("error reading llm response: %v", err)
880
881
	}

882
	return nil
883
884
}

885
886
type EmbedRequest struct {
	Content []string `json:"content"`
887
888
}

889
type EmbedResponse struct {
890
891
	Embedding       [][]float32 `json:"embedding"`
	PromptEvalCount int         `json:"prompt_n"`
892
893
}

894
func (s *llmServer) Embed(ctx context.Context, input []string) (*EmbedResponse, error) {
895
896
897
898
	// each input will use a slot, so we need to acquire the semaphore for
	// the number of inputs up to numParallel
	slots := int64(min(len(input), s.numParallel))
	if err := s.sem.Acquire(ctx, slots); err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
899
900
901
		slog.Error("Failed to acquire semaphore", "error", err)
		return nil, err
	}
902
	defer s.sem.Release(slots)
903

904
	// Make sure the server is ready
905
	status, err := s.getServerStatusRetry(ctx)
906
907
908
	if err != nil {
		return nil, err
	} else if status != ServerStatusReady {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
909
		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
910
911
	}

912
	data, err := json.Marshal(EmbedRequest{Content: input})
Michael Yang's avatar
Michael Yang committed
913
	if err != nil {
914
915
916
		return nil, fmt.Errorf("error marshaling embed data: %w", err)
	}

Michael Yang's avatar
Michael Yang committed
917
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/embedding", s.port), bytes.NewBuffer(data))
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
	if err != nil {
		return nil, fmt.Errorf("error creating embed request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	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 {
		log.Printf("llm encode error: %s", body)
		return nil, fmt.Errorf("%s", body)
	}

939
940
	var e EmbedResponse
	if err := json.Unmarshal(body, &e); err != nil {
941
942
943
		return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
	}

944
	return &e, nil
945
946
}

Michael Yang's avatar
Michael Yang committed
947
948
949
950
951
952
953
954
type TokenizeRequest struct {
	Content string `json:"content"`
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
955
func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error) {
Michael Yang's avatar
Michael Yang committed
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
	// 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()

	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"`
1005
1006
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1007
func (s *llmServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
Michael Yang's avatar
Michael Yang committed
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
	// 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()

	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
1049
1050
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1051
func (s *llmServer) Close() error {
1052
1053
	if s.cmd != nil {
		slog.Debug("stopping llama server")
1054
1055
1056
		if err := s.cmd.Process.Kill(); err != nil {
			return err
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1057
1058
1059
1060
1061
		// 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
		}
1062
1063

		slog.Debug("llama server stopped")
1064
1065
1066
1067
1068
	}

	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1069
func (s *llmServer) EstimatedVRAM() uint64 {
1070
	return s.estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1071
1072
}

1073
func (s *llmServer) EstimatedTotal() uint64 {
1074
	return s.estimate.TotalSize
1075
1076
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1077
func (s *llmServer) EstimatedVRAMByGPU(gpuID string) uint64 {
1078
1079
1080
1081
1082
1083
1084
1085
	for i, gpu := range s.gpus {
		if gpu.ID == gpuID {
			return s.estimate.GPUSizes[i]
		}
	}
	return 0
}

1086
1087
1088
1089
1090
1091
1092
1093
func parseDurationMs(ms float64) time.Duration {
	dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
	if err != nil {
		panic(err)
	}

	return dur
}