runner.go 21.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
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
	// TODO(jessegross): We should ensure that we always leave minBatch of context space to shift,
	// otherwise we might truncate or split the batch against the model's wishes

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

Jesse Gross's avatar
Jesse Gross committed
133
	// TODO(jessegross): Ingest cached history for grammar
134
135

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

160
	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
161

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

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

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

185
		for _, t := range tokens {
186
			inputs = append(inputs, input.Input{Token: t})
187
188
		}

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

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

212
213
214
215
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

216
			inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
217
218
219
220
221
222
			postTokenize = true
		}
	}

	if visionModel && postTokenize {
		var err error
223
		inputs, err = multimodalProcessor.PostTokenize(inputs)
224
		if err != nil {
225
			return nil, nil, err
226
227
228
		}
	}

229
	return inputs, &contexts, nil
230
231
232
}

type Server struct {
233
234
235
236
237
	// 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
238
	model model.Model
239

240
	// status for external health reporting - loading, ready to serve, etc.
241
	status llm.ServerStatus
242
243
244
245
246
247
248
249

	// 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)
250
	// TODO (jmorganca): make this n_batch
251
252
	batchSize int

253
254
255
256
257
258
259
260
	// 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
261
262
	seqs []*Sequence

263
264
265
266
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

267
268
269
	// KV cache
	cache *InputCache

270
271
272
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
273
274
275
276
277
278

	// 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
279
280
281
282
283
284
285
286
287
288
289
290
}

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

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

304
305
306
307
308
309
310
311
312
313
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
314
315
316
317
318
319
320
321
322
323
324
}

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
325
	s.seqsSem.Release(1)
326
327
328
329
330
331
332
333
334
335
}

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

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

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

351
	var options input.Options
352

353
	for i, seq := range s.seqs {
354
355
356
357
358
		if seq == nil {
			continue
		}

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

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

369
370
		batchSize := s.batchSize

371
		for j, inp := range seq.inputs {
372
373
374
375
376
377
378
379
380
			// 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 {
381
382
				break
			}
Jesse Gross's avatar
Jesse Gross committed
383

384
385
386
387
388
389
390
391
392
393
394
395
396
397
			// If the sum of our working set (already processed tokens, tokens we added to this
			// batch, required following tokens) exceeds the context size, then trigger a shift
			// now so we don't have to do one later when we can't break the batch.
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+minBatch) > s.cache.numCtx {
				if len(seq.pendingInputs) != 0 {
					break
				}

				err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
				if err != nil {
					return err
				}
			}

398
399
400
			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})
401
402
			}

Jesse Gross's avatar
Jesse Gross committed
403
404
405
406
			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)
407
			if j+1 == len(seq.inputs) {
Jesse Gross's avatar
Jesse Gross committed
408
409
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
410
			seq.pendingInputs = append(seq.pendingInputs, inp)
411
		}
412
413

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
414
415
	}

Jesse Gross's avatar
Jesse Gross committed
416
	if len(options.Inputs) == 0 {
417
		return nil
418
419
	}

Jesse Gross's avatar
Jesse Gross committed
420
421
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
422

Jesse Gross's avatar
Jesse Gross committed
423
	modelOutput, err := model.Forward(ctx, s.model, options)
424
	if err != nil {
425
		return fmt.Errorf("failed to decode batch: %w", err)
426
427
	}

428
	logits := modelOutput.Floats()
429

430
431
432
433
434
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
435
		// After calling Forward, pending inputs are now in the cache
436
437
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
438
			seq.pendingInputs = []input.Input{}
439
440
		}

441
442
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
443
444
445
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
446
447
448
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
449
450
		seq.numPredicted++
		if seq.numPredicted == 1 {
451
452
453
454
455
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
456
			// TODO(jessegross): Embedding support
457
458
459
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
460
461
462
		}

		// sample a token
463
464
465
		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
466
		if err != nil {
467
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
468
		}
469
470

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

485
		seq.inputs = []input.Input{{Token: token}}
486
487
488
489

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

Jesse Gross's avatar
Jesse Gross committed
490
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
491
492
493
494
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
495
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
			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]
511
512
513
514
515

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

Jesse Gross's avatar
Jesse Gross committed
516
		if common.ContainsStopSuffix(sequence, seq.stop) {
517
518
519
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
520
		if common.IncompleteUnicode(sequence) {
521
522
523
524
525
526
527
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
528
529

	return nil
530
531
532
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
533
	var req llm.CompletionRequest
534
535
536
537
538
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

539
540
541
542
543
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

544
545
546
547
548
549
550
551
552
553
	// 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
	}

554
555
556
557
558
559
560
561
562
563
	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
		}
	}

564
	sampler := sample.NewSampler(
565
566
567
568
569
		req.Options.Temperature,
		req.Options.TopK,
		req.Options.TopP,
		req.Options.MinP,
		req.Options.Seed,
570
		grammar,
571
572
	)

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

585
	// Ensure there is a place to put the sequence, released when removed from s.seqs
586
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
587
588
589
590
591
		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)
		}
592
593
594
		return
	}

595
	s.mu.Lock()
596
	found := false
597
598
	for i, sq := range s.seqs {
		if sq == nil {
599
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs)
600
601
602
603
604
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
605

606
607
			s.seqs[i] = seq
			s.cond.Signal()
608
			found = true
609
610
611
612
613
			break
		}
	}
	s.mu.Unlock()

614
615
616
617
618
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

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

668
669
670
671
672
673
674
675
676
677
678
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, ", ")
}

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

694
695
	s.vocab = sample.NewVocab(mpath)

Jesse Gross's avatar
Jesse Gross committed
696
	// TODO(jessegross): LoRA loading
697
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
698
		panic("loras are not yet implemented")
699
700
	}

Jesse Gross's avatar
Jesse Gross committed
701
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
702
703
704
	if err != nil {
		panic(err)
	}
705

Jesse Gross's avatar
Jesse Gross committed
706
707
708
709
710
711
712
713
714
	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))

715
	s.status = llm.ServerStatusReady
716
717
718
	s.ready.Done()
}

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

737
	var lpaths multiLPath
738
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
739

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

	server := &Server{
		batchSize: *batchSize,
767
		status:    llm.ServerStatusLoadingModel,
768
769
	}

Jesse Gross's avatar
Jesse Gross committed
770
771
772
773
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

774
	var tensorSplitFloats []float32
775
	if *tensorSplit != "" {
776
777
778
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
779
			f, _ := strconv.ParseFloat(s, 32)
780
			tensorSplitFloats[i] = float32(f)
781
		}
782
783
784
	}

	params := ml.BackendParams{
785
786
787
788
789
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
790
	}
791
792

	server.ready.Add(1)
793
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
794
795
796
797

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

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

800
801
802
803
804
805
	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)
806
		return err
807
808
809
810
	}
	defer listener.Close()

	mux := http.NewServeMux()
811
812
813
814
815
816
817
	// 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)
818
819
820
821
822
823
824
825

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

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

829
	return nil
830
}