runner.go 20.7 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
38
39
40
type contextList struct {
	list []ml.Context
}

41
type Sequence struct {
42
	// ctxs are used for allocating tensors that last the lifetime of the sequence, such as
43
	// multimodal embeddings
44
	ctxs *contextList
45

46
47
48
49
	// batch index
	iBatch int

	// prompt inputs left to evaluate
50
	inputs []input.Input
51

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

55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
	// 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

70
71
	// sampler with transforms to run on generated logits
	sampler sample.Sampler
72
73
74
75
76
77
78
79

	// 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
80
	numKeep int32
81
82
83
84
85
86
87
88
89

	// 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
90
	numPredicted        int
91
92
93
94
	numPromptInputs     int
}

type NewSequenceParams struct {
Jesse Gross's avatar
Jesse Gross committed
95
96
97
	numPredict int
	stop       []string
	numKeep    int32
98
	sampler    sample.Sampler
Jesse Gross's avatar
Jesse Gross committed
99
	embedding  bool
100
101
}

102
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
103
104
105
106
	s.ready.Wait()

	startTime := time.Now()

107
	inputs, ctxs, err := s.inputs(prompt, images)
108
109
110
111
112
113
114
	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
115
		params.numKeep = int32(len(inputs))
116
117
	}

118
119
120
	// 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
121
122
	if int32(len(inputs)) > s.cache.numCtx {
		discard := int32(len(inputs)) - s.cache.numCtx
123
		newInputs := inputs[:params.numKeep]
124
125
126
		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))
127
		inputs = newInputs
128
129
	}

Jesse Gross's avatar
Jesse Gross committed
130
	// TODO(jessegross): Ingest cached history for grammar
131
132

	return &Sequence{
133
		ctxs:                ctxs,
134
135
136
137
138
139
140
141
		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),
142
		sampler:             params.sampler,
143
144
145
146
147
148
149
150
		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
151
// decoding images
152
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input.Input, *contextList, error) {
153
	var inputs []input.Input
154
155
156
	var parts []string
	var matches [][]string

157
	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
158

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

167
168
169
170
171
172
173
	var contexts contextList
	runtime.AddCleanup(&contexts, func(ctxs []ml.Context) {
		for _, ctx := range ctxs {
			ctx.Close()
		}
	}, contexts.list)

174
	postTokenize := false
175
176
	for i, part := range parts {
		// text - tokenize
177
		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
178
		if err != nil {
179
			return nil, nil, err
180
		}
181

182
		for _, t := range tokens {
183
			inputs = append(inputs, input.Input{Token: t})
184
185
		}

Jesse Gross's avatar
Jesse Gross committed
186
		// image - decode and store
187
188
189
190
191
192
193
194
195
196
197
198
		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 {
199
				return nil, nil, fmt.Errorf("invalid image index: %d", n)
200
201
			}

202
203
			ctx := s.model.Backend().NewContext()
			contexts.list = append(contexts.list, ctx)
204
			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
205
			if err != nil {
206
				return nil, nil, err
Jesse Gross's avatar
Jesse Gross committed
207
208
			}

209
210
211
212
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

213
			inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
214
215
216
217
218
219
			postTokenize = true
		}
	}

	if visionModel && postTokenize {
		var err error
220
		inputs, err = multimodalProcessor.PostTokenize(inputs)
221
		if err != nil {
222
			return nil, nil, err
223
224
225
		}
	}

226
	return inputs, &contexts, nil
227
228
229
}

type Server struct {
230
231
232
233
234
	// 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
235
	model model.Model
236

237
	// status for external health reporting - loading, ready to serve, etc.
238
	status llm.ServerStatus
239
240
241
242
243
244
245
246

	// 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)
247
	// TODO (jmorganca): make this n_batch
248
249
	batchSize int

250
251
252
253
254
255
256
257
	// 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
258
259
	seqs []*Sequence

260
261
262
263
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

264
265
266
	// KV cache
	cache *InputCache

267
268
269
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
270
271
272
273
274
275

	// 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
276
277
278
279
280
281
282
283
284
285
286
287
}

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

func flushPending(seq *Sequence) bool {
288
289
290
291
292
293
294
295
296
297
298
	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]
299
300
	}

301
302
303
304
305
306
307
308
309
310
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
311
312
313
314
315
316
317
318
319
320
321
}

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
	s.seqs[seqIndex] = nil
322
	s.seqsSem.Release(1)
323
324
325
326
327
328
329
330
331
332
}

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

	for {
		select {
		case <-ctx.Done():
			return
		default:
Jesse Gross's avatar
Jesse Gross committed
333
			err := s.processBatch()
334
335
336
			if err != nil {
				panic(err)
			}
337
338
339
340
		}
	}
}

Jesse Gross's avatar
Jesse Gross committed
341
func (s *Server) processBatch() error {
342
343
344
345
346
347
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

348
	var options input.Options
349

350
	for i, seq := range s.seqs {
351
352
353
354
355
		if seq == nil {
			continue
		}

		// if past the num predict limit
356
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
357
			s.removeSequence(i, "limit")
358
359
360
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
361
362
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
363
			seq.cache.Inputs = []input.Input{}
Jesse Gross's avatar
Jesse Gross committed
364
365
		}

366
367
		batchSize := s.batchSize

368
		for j, inp := range seq.inputs {
Jesse Gross's avatar
Jesse Gross committed
369
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
370
371
372
373
374
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
						return err
					}
375
376
377
378
379
				} else {
					break
				}
			}

380
381
382
383
384
385
386
387
388
			// If we are required to put following inputs into a single batch then extend the
			// batch size. Since we are only extending the size the minimum amount possible, this
			// will cause a break if we have pending inputs.
			minBatch := 1 + inp.SameBatch
			if minBatch > batchSize {
				batchSize = minBatch
			}

			if len(seq.pendingInputs)+minBatch > batchSize {
389
390
				break
			}
Jesse Gross's avatar
Jesse Gross committed
391

392
393
394
			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})
395
396
			}

Jesse Gross's avatar
Jesse Gross committed
397
398
399
400
			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)
401
			if j+1 == len(seq.inputs) {
Jesse Gross's avatar
Jesse Gross committed
402
403
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
404
			seq.pendingInputs = append(seq.pendingInputs, inp)
405
		}
406
407

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
408
409
	}

Jesse Gross's avatar
Jesse Gross committed
410
	if len(options.Inputs) == 0 {
411
		return nil
412
413
	}

Jesse Gross's avatar
Jesse Gross committed
414
415
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
416

Jesse Gross's avatar
Jesse Gross committed
417
	modelOutput, err := model.Forward(ctx, s.model, options)
418
	if err != nil {
419
		return fmt.Errorf("failed to decode batch: %w", err)
420
421
	}

422
	logits := modelOutput.Floats()
423

424
425
426
427
428
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
429
		// After calling Forward, pending inputs are now in the cache
430
431
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
432
			seq.pendingInputs = []input.Input{}
433
434
		}

435
436
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
437
438
439
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
440
441
442
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
443
444
		seq.numPredicted++
		if seq.numPredicted == 1 {
445
446
447
448
449
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
450
			// TODO(jessegross): Embedding support
451
452
453
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
454
455
456
		}

		// sample a token
457
458
459
		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
460
		if err != nil {
461
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
462
		}
463
464

		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
465
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
466
467
468
469
470
471
472
473
			// 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
474
475
476
477
478
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
			return err
		}

479
		seq.inputs = []input.Input{{Token: token}}
480
481
482
483

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

Jesse Gross's avatar
Jesse Gross committed
484
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
485
486
487
488
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
489
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
			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]
505
506
507
508
509

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

Jesse Gross's avatar
Jesse Gross committed
510
		if common.ContainsStopSuffix(sequence, seq.stop) {
511
512
513
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
514
		if common.IncompleteUnicode(sequence) {
515
516
517
518
519
520
521
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
522
523

	return nil
524
525
526
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
527
	var req llm.CompletionRequest
528
529
530
531
532
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

533
534
535
536
537
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

538
539
540
541
542
543
544
545
546
547
	// 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
	}

548
549
550
551
552
553
554
555
556
557
	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
		}
	}

558
	sampler := sample.NewSampler(
559
560
561
562
563
		req.Options.Temperature,
		req.Options.TopK,
		req.Options.TopP,
		req.Options.MinP,
		req.Options.Seed,
564
		grammar,
565
566
	)

567
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
568
569
570
		numPredict: req.Options.NumPredict,
		stop:       req.Options.Stop,
		numKeep:    int32(req.Options.NumKeep),
571
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
572
		embedding:  false,
573
574
575
576
577
578
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

579
	// Ensure there is a place to put the sequence, released when removed from s.seqs
580
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
581
582
583
584
585
		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)
		}
586
587
588
		return
	}

589
	s.mu.Lock()
590
	found := false
591
592
	for i, sq := range s.seqs {
		if sq == nil {
593
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs)
594
595
596
597
598
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
599

600
601
			s.seqs[i] = seq
			s.cond.Signal()
602
			found = true
603
604
605
606
607
			break
		}
	}
	s.mu.Unlock()

608
609
610
611
612
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

613
614
615
616
617
618
619
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
620
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
621
622
623
624
625
626
627
628
629
630
					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
631
632
633
634
635
636
637
638
639
640
641
				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),
642
643
644
645
646
647
648
649
650
651
652
653
				}); 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")
654
655
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
656
657
658
659
660
661
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

662
663
664
665
666
667
668
669
670
671
672
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, ", ")
}

673
674
func (s *Server) loadModel(
	mpath string,
675
	params ml.BackendParams,
676
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
677
	parallel int,
678
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
679
	kvSize int,
680
681
	multiUserCache bool,
) {
682
	var err error
683
	s.model, err = model.New(mpath, params)
684
685
686
	if err != nil {
		panic(err)
	}
687

688
689
	s.vocab = sample.NewVocab(mpath)

Jesse Gross's avatar
Jesse Gross committed
690
	// TODO(jessegross): LoRA loading
691
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
692
		panic("loras are not yet implemented")
693
694
	}

Jesse Gross's avatar
Jesse Gross committed
695
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
696
697
698
	if err != nil {
		panic(err)
	}
699

Jesse Gross's avatar
Jesse Gross committed
700
701
702
703
704
705
706
707
708
	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))

709
	s.status = llm.ServerStatusReady
710
711
712
	s.ready.Done()
}

713
714
715
716
717
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")
718
719
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
720
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
721
722
723
	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")
724
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
725
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
726
727
	_ = 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")
728
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
729
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
730

731
	var lpaths multiLPath
732
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
733

734
735
736
737
738
739
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
	}
	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
757
	slog.Info("starting ollama engine")
758
759
760

	server := &Server{
		batchSize: *batchSize,
761
		status:    llm.ServerStatusLoadingModel,
762
763
	}

Jesse Gross's avatar
Jesse Gross committed
764
765
766
767
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

768
	var tensorSplitFloats []float32
769
	if *tensorSplit != "" {
770
771
772
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
773
			f, _ := strconv.ParseFloat(s, 32)
774
			tensorSplitFloats[i] = float32(f)
775
		}
776
777
778
	}

	params := ml.BackendParams{
779
780
781
782
783
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
784
	}
785
786

	server.ready.Add(1)
787
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
788
789
790
791

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

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

794
795
796
797
798
799
	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)
800
		return err
801
802
803
804
	}
	defer listener.Close()

	mux := http.NewServeMux()
805
806
807
808
809
810
811
	// 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)
812
813
814
815
816
817
818
819

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

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

823
	return nil
824
}