runner.go 20.1 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package ollamarunner
2
3
4
5
6
7
8

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
9
	"hash/maphash"
10
11
12
13
14
15
16
17
18
19
20
21
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
22
	"unicode/utf8"
23

24
25
	"golang.org/x/sync/semaphore"

26
	"github.com/ollama/ollama/api"
27
	"github.com/ollama/ollama/llm"
28
	"github.com/ollama/ollama/ml"
Jesse Gross's avatar
Jesse Gross committed
29
	"github.com/ollama/ollama/model"
30
	"github.com/ollama/ollama/model/input"
Jesse Gross's avatar
Jesse Gross committed
31
32
33
34
	"github.com/ollama/ollama/runner/common"
	"github.com/ollama/ollama/sample"

	_ "github.com/ollama/ollama/model/models"
35
36
37
)

type Sequence struct {
38
39
40
41
	// ctx for allocating tensors that last the lifetime of the sequence, such as
	// multimodal embeddings
	ctx ml.Context

42
43
44
45
	// batch index
	iBatch int

	// prompt inputs left to evaluate
46
	inputs []input.Input
47

Jesse Gross's avatar
Jesse Gross committed
48
	// inputs that have been added to a batch but not yet submitted to Forward
49
	pendingInputs []input.Input
50

51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
	// tokens that have been generated but not returned yet (e.g. for stop sequences)
	pendingResponses []string

	// input cache being used by this sequence
	cache *InputCacheSlot

	// channel to send responses over
	responses chan string

	// channel to stop decoding (such as if the remote connection is closed)
	quit chan bool

	// number of tokens to predict
	numPredict int

66
67
	// sampler with transforms to run on generated logits
	sampler sample.Sampler
68
69
70
71
72
73
74
75

	// channel to send back the embedding if embedding only
	embedding chan []float32

	// stop sequences
	stop []string

	// number of inputs to keep at the beginning when shifting context window
Jesse Gross's avatar
Jesse Gross committed
76
	numKeep int32
77
78
79
80
81
82
83
84
85

	// true if an embedding are to be returned instead of text generation
	embeddingOnly bool

	doneReason string

	// Metrics
	startProcessingTime time.Time
	startGenerationTime time.Time
Jesse Gross's avatar
Jesse Gross committed
86
	numPredicted        int
87
88
89
90
	numPromptInputs     int
}

type NewSequenceParams struct {
Jesse Gross's avatar
Jesse Gross committed
91
92
93
	numPredict int
	stop       []string
	numKeep    int32
94
	sampler    sample.Sampler
Jesse Gross's avatar
Jesse Gross committed
95
	embedding  bool
96
97
}

98
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
99
100
101
	s.ready.Wait()

	startTime := time.Now()
102
	ctx := s.model.Backend().NewContext()
103

104
	inputs, err := s.inputs(ctx, prompt, images)
105
106
107
108
109
110
111
	if err != nil {
		return nil, fmt.Errorf("failed to process inputs: %w", err)
	} else if len(inputs) == 0 {
		return nil, errors.New("no input provided")
	}

	if params.numKeep < 0 {
Jesse Gross's avatar
Jesse Gross committed
112
		params.numKeep = int32(len(inputs))
113
114
	}

115
116
117
	// Ensure that at least 1 input can be discarded during shift
	params.numKeep = min(params.numKeep, s.cache.numCtx-1)

Jesse Gross's avatar
Jesse Gross committed
118
119
	if int32(len(inputs)) > s.cache.numCtx {
		discard := int32(len(inputs)) - s.cache.numCtx
120
		newInputs := inputs[:params.numKeep]
121
122
123
		newInputs = append(newInputs, inputs[params.numKeep+discard:]...)

		slog.Warn("truncating input prompt", "limit", s.cache.numCtx, "prompt", len(inputs), "keep", params.numKeep, "new", len(newInputs))
124
		inputs = newInputs
125
126
	}

Jesse Gross's avatar
Jesse Gross committed
127
	// TODO(jessegross): Ingest cached history for grammar
128
129

	return &Sequence{
130
		ctx:                 ctx,
131
132
133
134
135
136
137
138
		inputs:              inputs,
		numPromptInputs:     len(inputs),
		startProcessingTime: startTime,
		numPredict:          params.numPredict,
		pendingResponses:    make([]string, 0),
		responses:           make(chan string, 100),
		quit:                make(chan bool, 1),
		embedding:           make(chan []float32, 1),
139
		sampler:             params.sampler,
140
141
142
143
144
145
146
147
		embeddingOnly:       params.embedding,
		stop:                params.stop,
		numKeep:             params.numKeep,
	}, nil
}

// inputs processes the prompt and images into a list of inputs
// by splitting the prompt on [img-<n>] tags, tokenizing text and
Jesse Gross's avatar
Jesse Gross committed
148
// decoding images
149
func (s *Server) inputs(ctx ml.Context, prompt string, images []llm.ImageData) ([]input.Input, error) {
150
	var inputs []input.Input
151
152
153
	var parts []string
	var matches [][]string

154
	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
155

156
157
158
159
160
161
162
163
164
	if visionModel {
		re := regexp.MustCompile(`\[img-(\d+)\]`)
		parts = re.Split(prompt, -1)
		matches = re.FindAllStringSubmatch(prompt, -1)
	} else {
		parts = []string{prompt}
	}

	postTokenize := false
165
166
	for i, part := range parts {
		// text - tokenize
167
		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
168
169
170
		if err != nil {
			return nil, err
		}
171

172
		for _, t := range tokens {
173
			inputs = append(inputs, input.Input{Token: t})
174
175
		}

Jesse Gross's avatar
Jesse Gross committed
176
		// image - decode and store
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
		if i < len(matches) {
			n, _ := strconv.Atoi(matches[i][1])

			imageIndex := -1
			for j := range images {
				if images[j].ID == n {
					imageIndex = j
					break
				}
			}

			if imageIndex < 0 {
				return nil, fmt.Errorf("invalid image index: %d", n)
			}

192
			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
193
194
195
196
			if err != nil {
				return nil, err
			}

197
198
199
200
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

201
			inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
202
203
204
205
206
207
208
209
210
			postTokenize = true
		}
	}

	if visionModel && postTokenize {
		var err error
		inputs, err = multimodalProcessor.PostTokenize(ctx, inputs)
		if err != nil {
			return nil, err
211
212
213
214
215
216
217
		}
	}

	return inputs, nil
}

type Server struct {
218
219
220
221
222
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
Jesse Gross's avatar
Jesse Gross committed
223
	model model.Model
224

225
	// status for external health reporting - loading, ready to serve, etc.
226
	status llm.ServerStatus
227
228
229
230
231
232
233
234

	// current progress on loading the model
	progress float32

	// number of simultaneous requests to handle
	parallel int

	// maximum number of elements in a batch (per sequence)
235
	// TODO (jmorganca): make this n_batch
236
237
	batchSize int

238
239
240
241
242
243
244
245
	// protects access to everything below this line
	// this is context state needed for decoding
	mu sync.Mutex

	// indicates that data is ready for processing
	cond *sync.Cond

	// the list of simultaneous sequences being evaluated
246
247
	seqs []*Sequence

248
249
250
251
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

252
253
254
	// KV cache
	cache *InputCache

255
256
257
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
258
259
260
261
262
263

	// vocab is a llama.cpp vocab required for gammar-based
	// constrained generation (json mode, structured outputs)
	// TODO: this is temporary until Ollama sampling supports
	// constrained generation
	vocab *sample.Vocab
264
265
266
267
268
269
270
271
272
273
274
275
}

func (s *Server) allNil() bool {
	for _, item := range s.seqs {
		if item != nil {
			return false
		}
	}
	return true
}

func flushPending(seq *Sequence) bool {
276
277
278
279
280
281
282
283
284
285
286
	joined := strings.Join(seq.pendingResponses, "")
	seq.pendingResponses = []string{}

	// Check if there are any partial UTF-8 characters remaining.
	// We already check and queue as we are generating but some may
	// still make it here:
	// - Sequence is ending, e.g. generation limit has been hit
	// - Invalid characters in the middle of a string
	// This is a stricter check to ensure we never output invalid Unicode.
	for !utf8.ValidString(joined) {
		joined = joined[:len(joined)-1]
287
288
	}

289
290
291
292
293
294
295
296
297
298
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
299
300
301
302
303
304
305
306
307
308
}

func (s *Server) removeSequence(seqIndex int, reason string) {
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
309
	seq.ctx.Close()
310
	s.seqs[seqIndex] = nil
311
	s.seqsSem.Release(1)
312
313
314
315
316
317
318
319
320
321
}

func (s *Server) run(ctx context.Context) {
	s.ready.Wait()

	for {
		select {
		case <-ctx.Done():
			return
		default:
Jesse Gross's avatar
Jesse Gross committed
322
			err := s.processBatch()
323
324
325
			if err != nil {
				panic(err)
			}
326
327
328
329
		}
	}
}

Jesse Gross's avatar
Jesse Gross committed
330
func (s *Server) processBatch() error {
331
332
333
334
335
336
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

337
	var options input.Options
338

339
	for i, seq := range s.seqs {
340
341
342
343
344
		if seq == nil {
			continue
		}

		// if past the num predict limit
345
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
346
			s.removeSequence(i, "limit")
347
348
349
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
350
351
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
352
			seq.cache.Inputs = []input.Input{}
Jesse Gross's avatar
Jesse Gross committed
353
354
		}

355
		for j, inp := range seq.inputs {
Jesse Gross's avatar
Jesse Gross committed
356
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
357
358
359
360
361
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
						return err
					}
362
363
364
365
366
				} else {
					break
				}
			}

367
			if j >= s.batchSize {
368
369
				break
			}
Jesse Gross's avatar
Jesse Gross committed
370

371
372
373
			options.Inputs = append(options.Inputs, inp.Token)
			if inp.Multimodal != nil {
				options.Multimodal = append(options.Multimodal, input.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: inp.Multimodal})
374
375
			}

Jesse Gross's avatar
Jesse Gross committed
376
377
378
379
			options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
			options.Sequences = append(options.Sequences, seq.cache.Id)

			seq.iBatch = len(options.Outputs)
380
			if j+1 == len(seq.inputs) {
Jesse Gross's avatar
Jesse Gross committed
381
382
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
383
			seq.pendingInputs = append(seq.pendingInputs, inp)
384
		}
385
386

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
387
388
	}

Jesse Gross's avatar
Jesse Gross committed
389
	if len(options.Inputs) == 0 {
390
		return nil
391
392
	}

Jesse Gross's avatar
Jesse Gross committed
393
394
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
395

Jesse Gross's avatar
Jesse Gross committed
396
	modelOutput, err := model.Forward(ctx, s.model, options)
397
	if err != nil {
398
		return fmt.Errorf("failed to decode batch: %w", err)
399
400
	}

401
	logits := modelOutput.Floats()
402

403
404
405
406
407
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
408
		// After calling Forward, pending inputs are now in the cache
409
410
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
411
			seq.pendingInputs = []input.Input{}
412
413
		}

414
415
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
416
417
418
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
419
420
421
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
422
423
		seq.numPredicted++
		if seq.numPredicted == 1 {
424
425
426
427
428
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
429
			// TODO(jessegross): Embedding support
430
431
432
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
433
434
435
		}

		// sample a token
436
437
438
		vocabSize := len(logits) / len(options.Outputs)

		token, err := seq.sampler.Sample(logits[seq.iBatch*vocabSize : (seq.iBatch+1)*vocabSize])
Jesse Gross's avatar
Jesse Gross committed
439
		if err != nil {
440
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
441
		}
442
443

		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
444
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
445
446
447
448
449
450
451
452
			// TODO (jmorganca): we should send this back
			// as it's important for the /api/generate context
			// seq.responses <- piece

			s.removeSequence(i, "stop")
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
453
454
455
456
457
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
			return err
		}

458
		seq.inputs = []input.Input{{Token: token}}
459
460
461
462

		seq.pendingResponses = append(seq.pendingResponses, piece)
		sequence := strings.Join(seq.pendingResponses, "")

Jesse Gross's avatar
Jesse Gross committed
463
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
464
465
466
467
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
468
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
			newLen := len(seq.pendingResponses)

			// Update the cache based on the tokens that will be returned:
			// - We have 1 token more than is currently in the cache because
			// the last one generated wasn't submitted to Decode
			// - Remove any stop sequences that we stripped out
			// - If truncateStop removed a portion of a token, drop that
			// - As defense-in-depth, if truncatedToken didn't find a stop token
			// remove the extra one that we added to the cache len
			tokenLen := len(seq.cache.Inputs) + 1
			tokenLen -= origLen - newLen
			if tokenTruncated || origLen == newLen {
				tokenLen--
			}
			seq.cache.Inputs = seq.cache.Inputs[:tokenLen]
484
485
486
487
488

			s.removeSequence(i, "stop")
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
489
		if common.ContainsStopSuffix(sequence, seq.stop) {
490
491
492
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
493
		if common.IncompleteUnicode(sequence) {
494
495
496
497
498
499
500
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
501
502

	return nil
503
504
505
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
506
	var req llm.CompletionRequest
507
508
509
510
511
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

512
513
514
515
516
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

517
518
519
520
521
522
523
524
525
526
	// Set the headers to indicate streaming
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Transfer-Encoding", "chunked")

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", http.StatusInternalServerError)
		return
	}

527
528
529
530
531
532
533
534
535
536
	var grammar *sample.Grammar
	var err error
	if req.Grammar != "" {
		grammar, err = sample.NewGrammar(s.vocab, req.Grammar)
		if err != nil {
			http.Error(w, "failed to load model vocabulary required for format", http.StatusInternalServerError)
			return
		}
	}

537
	sampler := sample.NewSampler(
538
539
540
541
542
		req.Options.Temperature,
		req.Options.TopK,
		req.Options.TopP,
		req.Options.MinP,
		req.Options.Seed,
543
		grammar,
544
545
	)

546
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
547
548
549
		numPredict: req.Options.NumPredict,
		stop:       req.Options.Stop,
		numKeep:    int32(req.Options.NumKeep),
550
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
551
		embedding:  false,
552
553
554
555
556
557
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

558
	// Ensure there is a place to put the sequence, released when removed from s.seqs
559
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
560
561
562
563
564
		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)
		}
565
566
567
		return
	}

568
	s.mu.Lock()
569
	found := false
570
571
	for i, sq := range s.seqs {
		if sq == nil {
572
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
573
574
575
576
577
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
578

579
580
			s.seqs[i] = seq
			s.cond.Signal()
581
			found = true
582
583
584
585
586
			break
		}
	}
	s.mu.Unlock()

587
588
589
590
591
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

592
593
594
595
596
597
598
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
599
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
600
601
602
603
604
605
606
607
608
609
					Content: content,
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
				// Send the final response
610
611
612
613
614
615
616
617
618
619
620
				doneReason := "stop"
				if seq.doneReason == "limit" {
					doneReason = "length"
				}
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
					DoneReason:         doneReason,
					PromptEvalCount:    seq.numPromptInputs,
					PromptEvalDuration: seq.startGenerationTime.Sub(seq.startProcessingTime),
					EvalCount:          seq.numPredicted,
					EvalDuration:       time.Since(seq.startGenerationTime),
621
622
623
624
625
626
627
628
629
630
631
632
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
				}

				return
			}
		}
	}
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
633
634
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
635
636
637
638
639
640
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

641
642
643
644
645
646
647
648
649
650
651
type multiLPath []string

func (m *multiLPath) Set(value string) error {
	*m = append(*m, value)
	return nil
}

func (m *multiLPath) String() string {
	return strings.Join(*m, ", ")
}

652
653
func (s *Server) loadModel(
	mpath string,
654
	params ml.BackendParams,
655
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
656
	parallel int,
657
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
658
	kvSize int,
659
660
	multiUserCache bool,
) {
661
	var err error
662
	s.model, err = model.New(mpath, params)
663
664
665
	if err != nil {
		panic(err)
	}
666

667
668
	s.vocab = sample.NewVocab(mpath)

Jesse Gross's avatar
Jesse Gross committed
669
	// TODO(jessegross): LoRA loading
670
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
671
		panic("loras are not yet implemented")
672
673
	}

Jesse Gross's avatar
Jesse Gross committed
674
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
675
676
677
	if err != nil {
		panic(err)
	}
678

Jesse Gross's avatar
Jesse Gross committed
679
680
681
682
683
684
685
686
687
	if !s.cache.enabled && parallel > 1 {
		parallel = 1
		slog.Warn("model does not support caching, disabling parallel processing")
	}

	s.parallel = parallel
	s.seqs = make([]*Sequence, s.parallel)
	s.seqsSem = semaphore.NewWeighted(int64(s.parallel))

688
	s.status = llm.ServerStatusReady
689
690
691
	s.ready.Done()
}

692
693
694
695
696
func Execute(args []string) error {
	fs := flag.NewFlagSet("runner", flag.ExitOnError)
	mpath := fs.String("model", "", "Path to model binary file")
	parallel := fs.Int("parallel", 1, "Number of sequences to handle simultaneously")
	batchSize := fs.Int("batch-size", 512, "Batch size")
697
698
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
699
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
700
701
702
	kvSize := fs.Int("ctx-size", 2048, "Context (or KV cache) size")
	kvCacheType := fs.String("kv-cache-type", "", "quantization type for KV cache (default: f16)")
	port := fs.Int("port", 8080, "Port to expose the server on")
703
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
704
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
705
706
	_ = fs.Bool("no-mmap", false, "do not memory-map model (slower load but may reduce pageouts if not using mlock)")
	_ = fs.Bool("mlock", false, "force system to keep model in RAM rather than swapping or compressing")
707
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
708
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
709

710
	var lpaths multiLPath
711
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
712

713
714
715
716
717
718
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
	}
	level := slog.LevelInfo
	if *verbose {
		level = slog.LevelDebug
	}
	handler := slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level:     level,
		AddSource: true,
		ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
			if attr.Key == slog.SourceKey {
				source := attr.Value.Any().(*slog.Source)
				source.File = filepath.Base(source.File)
			}
			return attr
		},
	})
	slog.SetDefault(slog.New(handler))
Jesse Gross's avatar
Jesse Gross committed
736
	slog.Info("starting ollama engine")
737
738
739

	server := &Server{
		batchSize: *batchSize,
740
		status:    llm.ServerStatusLoadingModel,
741
742
	}

Jesse Gross's avatar
Jesse Gross committed
743
744
745
746
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

747
	var tensorSplitFloats []float32
748
	if *tensorSplit != "" {
749
750
751
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
752
			f, _ := strconv.ParseFloat(s, 32)
753
			tensorSplitFloats[i] = float32(f)
754
		}
755
756
757
	}

	params := ml.BackendParams{
758
759
760
761
762
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
763
	}
764
765

	server.ready.Add(1)
766
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
767
768
769
770

	server.cond = sync.NewCond(&server.mu)

	ctx, cancel := context.WithCancel(context.Background())
Michael Yang's avatar
Michael Yang committed
771
772
	defer cancel()

773
774
775
776
777
778
	go server.run(ctx)

	addr := "127.0.0.1:" + strconv.Itoa(*port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Println("Listen error:", err)
779
		return err
780
781
782
783
	}
	defer listener.Close()

	mux := http.NewServeMux()
784
785
786
787
788
789
790
	// TODO: support embeddings
	mux.HandleFunc("POST /embedding", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
	})

	mux.HandleFunc("POST /completion", server.completion)
	mux.HandleFunc("GET /health", server.health)
791
792
793
794
795
796
797
798

	httpServer := http.Server{
		Handler: mux,
	}

	log.Println("Server listening on", addr)
	if err := httpServer.Serve(listener); err != nil {
		log.Fatal("server error:", err)
799
		return err
800
801
	}

802
	return nil
803
}