runner.go 22 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
	// next sequence for prompt processing to avoid starvation
	nextSeq int

273
274
275
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
276
277
278
279
280
281

	// 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
282
283
284
285
286
287
288
289
290
291
292
293
}

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

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

307
308
309
310
311
312
313
314
315
316
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
317
318
319
320
321
322
323
324
325
326
327
}

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

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

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

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

354
	var batchInputs []int32
Jesse Gross's avatar
Jesse Gross committed
355
	var batch input.Batch
356

357
358
359
360
361
362
	resumeSeq := -1
	seqIdx := s.nextSeq - 1
	for range s.seqs {
		seqIdx = (seqIdx + 1) % len(s.seqs)
		seq := s.seqs[seqIdx]

363
364
365
366
367
		if seq == nil {
			continue
		}

		// if past the num predict limit
368
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
369
			s.removeSequence(seqIdx, "limit")
370
371
372
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
373
374
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
375
			seq.cache.Inputs = []input.Input{}
Jesse Gross's avatar
Jesse Gross committed
376
377
		}

378
379
		batchSize := s.batchSize

380
		for i, inp := range seq.inputs {
381
382
			// 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
383
			// will cause a break if we have existing inputs.
384
385
386
387
388
			minBatch := 1 + inp.SameBatch
			if minBatch > batchSize {
				batchSize = minBatch
			}

389
390
391
392
393
394
395
396
			// Stop if the required batch would put us over the total batch size (including tokens
			// added by other sequences). If we haven't been able to add anything yet then pick up
			// here again for the next batch to avoid starvation, though we can opportunistically
			// check if other sequences can still squeeze something in.
			if len(batchInputs)+minBatch > batchSize {
				if len(seq.pendingInputs) == 0 && resumeSeq == -1 {
					resumeSeq = seqIdx
				}
397
398
				break
			}
Jesse Gross's avatar
Jesse Gross committed
399

400
401
402
403
404
405
406
407
408
409
410
411
412
413
			// 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
				}
			}

414
			batchInputs = append(batchInputs, inp.Token)
415
			if inp.Multimodal != nil {
416
				batch.Multimodal = append(batch.Multimodal, input.MultimodalIndex{Index: len(batchInputs) - 1, Multimodal: inp.Multimodal})
417
418
			}

Jesse Gross's avatar
Jesse Gross committed
419
420
			batch.Positions = append(batch.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
			batch.Sequences = append(batch.Sequences, seq.cache.Id)
Jesse Gross's avatar
Jesse Gross committed
421

Jesse Gross's avatar
Jesse Gross committed
422
			seq.iBatch = len(batch.Outputs)
423
			if i+1 == len(seq.inputs) {
424
				batch.Outputs = append(batch.Outputs, int32(len(batchInputs)-1))
Jesse Gross's avatar
Jesse Gross committed
425
			}
426
			seq.pendingInputs = append(seq.pendingInputs, inp)
427
		}
428
429

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
430
431
	}

432
433
434
435
436
437
	if resumeSeq != -1 {
		s.nextSeq = resumeSeq
	} else {
		s.nextSeq = seqIdx + 1
	}

438
	if len(batchInputs) == 0 {
439
		return nil
440
441
	}

Jesse Gross's avatar
Jesse Gross committed
442
443
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
444

445
	modelOutput, err := model.Forward(ctx, s.model, batchInputs, batch)
446
	if err != nil {
447
		return fmt.Errorf("failed to decode batch: %w", err)
448
449
	}

450
	logits := modelOutput.Floats()
451

452
453
454
455
456
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
457
		// After calling Forward, pending inputs are now in the cache
458
459
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
460
			seq.pendingInputs = []input.Input{}
461
462
		}

463
464
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
465
466
467
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
468
469
470
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
471
472
		seq.numPredicted++
		if seq.numPredicted == 1 {
473
474
475
476
477
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
478
			// TODO(jessegross): Embedding support
479
480
481
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
482
483
484
		}

		// sample a token
Jesse Gross's avatar
Jesse Gross committed
485
		vocabSize := len(logits) / len(batch.Outputs)
486
487

		token, err := seq.sampler.Sample(logits[seq.iBatch*vocabSize : (seq.iBatch+1)*vocabSize])
Jesse Gross's avatar
Jesse Gross committed
488
		if err != nil {
489
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
490
		}
491
492

		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
493
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
494
495
496
497
498
499
500
501
			// 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
502
503
504
505
506
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
			return err
		}

507
		seq.inputs = []input.Input{{Token: token}}
508
509
510
511

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

Jesse Gross's avatar
Jesse Gross committed
512
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
513
514
515
516
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
517
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
			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]
533
534
535
536
537

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

Jesse Gross's avatar
Jesse Gross committed
538
		if common.ContainsStopSuffix(sequence, seq.stop) {
539
540
541
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
542
		if common.IncompleteUnicode(sequence) {
543
544
545
546
547
548
549
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
550
551

	return nil
552
553
554
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
555
	var req llm.CompletionRequest
556
557
558
559
560
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

561
562
563
564
565
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

566
567
568
569
570
571
572
573
574
575
	// 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
	}

576
577
578
579
580
581
582
583
584
585
	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
		}
	}

586
	sampler := sample.NewSampler(
587
588
589
590
591
		req.Options.Temperature,
		req.Options.TopK,
		req.Options.TopP,
		req.Options.MinP,
		req.Options.Seed,
592
		grammar,
593
594
	)

595
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
596
597
598
		numPredict: req.Options.NumPredict,
		stop:       req.Options.Stop,
		numKeep:    int32(req.Options.NumKeep),
599
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
600
		embedding:  false,
601
602
603
604
605
606
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

607
	// Ensure there is a place to put the sequence, released when removed from s.seqs
608
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
609
610
611
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
612
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
613
		}
614
615
616
		return
	}

617
	s.mu.Lock()
618
	found := false
619
620
	for i, sq := range s.seqs {
		if sq == nil {
621
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs)
622
623
			if err != nil {
				s.mu.Unlock()
624
				s.seqsSem.Release(1)
625
626
627
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
628

629
630
			s.seqs[i] = seq
			s.cond.Signal()
631
			found = true
632
633
634
635
636
			break
		}
	}
	s.mu.Unlock()

637
	if !found {
638
		s.seqsSem.Release(1)
639
640
641
642
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

643
644
645
646
647
648
649
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
650
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
651
652
653
654
655
656
657
658
659
660
					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
661
662
663
664
665
666
667
668
669
670
671
				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),
672
673
674
675
676
677
678
679
680
681
682
683
				}); 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")
684
685
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
686
687
688
689
690
691
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

692
693
694
695
696
697
698
699
700
701
702
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, ", ")
}

703
func (s *Server) loadModel(
704
	ctx context.Context,
705
	mpath string,
706
	params ml.BackendParams,
707
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
708
	parallel int,
709
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
710
	kvSize int,
711
712
	multiUserCache bool,
) {
713
	var err error
714
	s.model, err = model.New(ctx, mpath, params)
715
716
717
	if err != nil {
		panic(err)
	}
718

719
720
	s.vocab = sample.NewVocab(mpath)

Jesse Gross's avatar
Jesse Gross committed
721
	// TODO(jessegross): LoRA loading
722
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
723
		panic("loras are not yet implemented")
724
725
	}

726
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, s.batchSize, multiUserCache)
727
728
729
	if err != nil {
		panic(err)
	}
730

Jesse Gross's avatar
Jesse Gross committed
731
732
733
734
735
736
737
738
739
	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))

740
	s.status = llm.ServerStatusReady
741
742
743
	s.ready.Done()
}

744
745
746
747
748
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")
749
750
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
751
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
752
753
754
	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")
755
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
756
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
757
758
	_ = 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")
759
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
760
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
761

762
	var lpaths multiLPath
763
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
764

765
766
767
768
769
770
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
	}
	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
788
	slog.Info("starting ollama engine")
789
790
791

	server := &Server{
		batchSize: *batchSize,
792
		status:    llm.ServerStatusLoadingModel,
793
794
	}

Jesse Gross's avatar
Jesse Gross committed
795
796
797
798
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

799
	var tensorSplitFloats []float32
800
	if *tensorSplit != "" {
801
802
803
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
804
			f, _ := strconv.ParseFloat(s, 32)
805
			tensorSplitFloats[i] = float32(f)
806
		}
807
808
809
	}

	params := ml.BackendParams{
810
811
812
		Progress: func(progress float32) {
			server.progress = progress
		},
813
814
815
816
817
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
818
	}
819
820
821

	server.ready.Add(1)
	ctx, cancel := context.WithCancel(context.Background())
Michael Yang's avatar
Michael Yang committed
822
823
	defer cancel()

824
825
826
827
	go server.loadModel(ctx, *mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)

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

828
829
830
831
832
833
	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)
834
		return err
835
836
837
838
	}
	defer listener.Close()

	mux := http.NewServeMux()
839
840
841
842
843
844
845
	// 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)
846
847
848
849
850
851
852
853

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

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

857
	return nil
858
}