runner.go 25.2 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package llamarunner
2
3
4
5
6
7
8
9
10
11
12
13
14

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"regexp"
15
	"sort"
16
17
18
19
	"strconv"
	"strings"
	"sync"
	"time"
20
	"unicode/utf8"
21

22
23
	"golang.org/x/sync/semaphore"

24
	"github.com/ollama/ollama/api"
25
	"github.com/ollama/ollama/envconfig"
26
	"github.com/ollama/ollama/llama"
27
	"github.com/ollama/ollama/llm"
28
	"github.com/ollama/ollama/logutil"
29
	"github.com/ollama/ollama/ml"
Jesse Gross's avatar
Jesse Gross committed
30
	"github.com/ollama/ollama/runner/common"
31
32
)

33
34
35
36
37
38
// response contains a piece of generated text along with optional logprobs
type response struct {
	content  string
	logprobs []llm.Logprob
}

39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
// input is an element of the prompt to process, either
// a token or an image embedding (generated from a vision projector)
type input struct {
	token int

	// embed is an image embedding
	embed []float32
}

type Sequence struct {
	// batch index
	iBatch int

	// number of tokens predicted so far
	numPredicted int

	// prompt inputs left to evaluate
	inputs []input

58
59
60
	// inputs that have been added to a batch but not yet submitted to Decode
	pendingInputs []input

61
62
63
	// tokens that have been generated but not returned yet (e.g. for stop sequences)
	pendingResponses []string

64
65
66
	// logprobs for tokens that haven't been returned yet
	pendingLogprobs []llm.Logprob

67
68
69
70
	// input cache being used by this sequence
	cache *InputCacheSlot

	// channel to send responses over
71
	responses chan response
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

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

	// number of tokens to predict
	numPredict int

	samplingCtx *llama.SamplingContext

	// 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
	numKeep int

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

93
94
95
	// shift if context window is exceeded
	shift bool

96
	doneReason llm.DoneReason
97

98
99
100
101
	// logprobs configuration
	logprobs    bool
	topLogprobs int

102
	// Metrics
Michael Yang's avatar
Michael Yang committed
103
104
105
106
	processingDuration time.Duration
	generationDuration time.Duration
	numDecoded         int
	numPromptInputs    int
107
108
109
110
111
112
113
114
}

type NewSequenceParams struct {
	numPredict     int
	stop           []string
	numKeep        int
	samplingParams *llama.SamplingParams
	embedding      bool
115
116
	shift          bool
	truncate       bool
117
118
	logprobs       bool
	topLogprobs    int
119
120
}

121
122
var errorInputTooLong = errors.New("the input length exceeds the context length")

123
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
124
125
126
127
128
129
130
131
132
133
134
135
136
	s.ready.Wait()

	inputs, err := s.inputs(prompt, images)
	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 {
		params.numKeep = len(inputs)
	}

137
138
	if s.model.AddBOSToken() {
		params.numKeep += 1
139
140
	}

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

144
	if len(inputs) > s.cache.numCtx {
145
		discard := len(inputs) - s.cache.numCtx
146
147
148
149
		if !params.truncate {
			return nil, errorInputTooLong
		}

150
		newInputs := inputs[:params.numKeep]
151
152
153
		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))
154
		inputs = newInputs
155
156
157
158
	}

	var sc *llama.SamplingContext
	if params.samplingParams != nil {
Jesse Gross's avatar
Jesse Gross committed
159
160
161
162
		sc, err = llama.NewSamplingContext(s.model, *params.samplingParams)
		if err != nil {
			return nil, err
		}
163
164
		for _, input := range inputs {
			if input.embed == nil {
165
				sc.Accept(input.token, false)
166
167
168
169
170
			}
		}
	}

	return &Sequence{
Michael Yang's avatar
Michael Yang committed
171
172
173
174
		inputs:           inputs,
		numPromptInputs:  len(inputs),
		numPredict:       params.numPredict,
		pendingResponses: make([]string, 0),
175
		responses:        make(chan response, 100),
Michael Yang's avatar
Michael Yang committed
176
177
178
179
180
181
		quit:             make(chan bool, 1),
		embedding:        make(chan []float32, 1),
		samplingCtx:      sc,
		embeddingOnly:    params.embedding,
		stop:             params.stop,
		numKeep:          params.numKeep,
182
		shift:            params.shift,
183
184
		logprobs:         params.logprobs,
		topLogprobs:      params.topLogprobs,
185
186
187
	}, nil
}

188
189
190
191
192
// calculateLogprobsLlama converts raw logits to log probabilities and finds top K tokens
func calculateLogprobsLlama(logits []float32, selectedToken int, topK int, model *llama.Model) []llm.Logprob {
	return common.CalculateLogprobs(logits, selectedToken, topK, model.TokenToPiece)
}

193
194
195
// inputs processes the prompt and images into a list of inputs
// by splitting the prompt on [img-<n>] tags, tokenizing text and
// generating image embeddings for each image
196
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
197
	var inputs []input
198
199
200
201
202
203
204
205
206
207
	var parts []string
	var matches [][]string

	if s.image != nil {
		re := regexp.MustCompile(`\[img-(\d+)\]`)
		parts = re.Split(prompt, -1)
		matches = re.FindAllStringSubmatch(prompt, -1)
	} else {
		parts = []string{prompt}
	}
208
209
210

	for i, part := range parts {
		// text - tokenize
211
212
213
214
		tokens, err := s.lc.Model().Tokenize(part, i == 0, true)
		if err != nil {
			return nil, err
		}
215

216
217
		for _, t := range tokens {
			inputs = append(inputs, input{token: t})
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
		}

		// image - generate image embedding
		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)
			}

236
			chunks, err := s.image.MultimodalTokenize(s.lc, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
237
238
239
240
			if err != nil {
				return nil, err
			}

241
242
243
244
245
246
247
248
			for _, c := range chunks {
				if len(c.Embed) != 0 {
					inputs = append(inputs, input{embed: c.Embed})
				} else {
					for _, t := range c.Tokens {
						inputs = append(inputs, input{token: t})
					}
				}
249
250
251
252
253
254
255
256
			}
		}
	}

	return inputs, nil
}

type Server struct {
Jesse Gross's avatar
Jesse Gross committed
257
258
259
260
261
262
	// modelPath is the location of the model to be loaded
	modelPath string

	// loadMu prevents more than one load attempt from occurring at a time
	loadMu sync.Mutex

263
264
265
266
267
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
268
269
	model *llama.Model

270
	// image model context for multi-modal models
271
	image *ImageContext
272

273
	// status for external health reporting - loading, ready to serve, etc.
274
	status llm.ServerStatus
275
276
277
278
279
280
281
282

	// 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)
283
	// TODO (jmorganca): make this n_batch
284
285
	batchSize int

286
287
288
289
290
291
292
293
294
	// 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

	// decoding state
	lc *llama.Context
295

296
	// the list of simultaneous sequences being evaluated
297
298
	seqs []*Sequence

299
300
301
302
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
	// KV cache
	cache *InputCache

	// next sequence for prompt processing to avoid starvation
	nextSeq int
}

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

func flushPending(seq *Sequence) bool {
320
	joined := strings.Join(seq.pendingResponses, "")
321
	logprobs := seq.pendingLogprobs
322
	seq.pendingResponses = []string{}
323
	seq.pendingLogprobs = []llm.Logprob{}
324
325
326
327
328
329
330
331
332

	// 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]
333
334
	}

335
336
337
338
339
	if len(joined) == 0 {
		return true
	}

	select {
340
	case seq.responses <- response{content: joined, logprobs: logprobs}:
341
342
343
344
		return true
	case <-seq.quit:
		return false
	}
345
346
}

347
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
348
349
350
351
352
353
354
355
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
356
	s.seqsSem.Release(1)
357
358
359
360
361
}

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

362
	// Logically these batches are used only within the context of processBatch
363
	// but it is better for performance to allocate them once here
Jesse Gross's avatar
Jesse Gross committed
364
365
366
367
	tokenBatch, err := llama.NewBatch(s.batchSize, len(s.seqs), 0)
	if err != nil {
		panic(err)
	}
368
369
	defer tokenBatch.Free()

370
371
372
	var embedBatch *llama.Batch
	embedBatchSize := s.image.BatchSize(s.batchSize)
	if embedBatchSize != 0 {
Jesse Gross's avatar
Jesse Gross committed
373
374
375
376
		embedBatch, err = llama.NewBatch(embedBatchSize, len(s.seqs), s.image.EmbedSize(s.lc))
		if err != nil {
			panic(err)
		}
377
378
379
380
		defer embedBatch.Free()
	} else {
		embedBatch = &llama.Batch{}
	}
381
382
383
384
385
386

	for {
		select {
		case <-ctx.Done():
			return
		default:
387
388
389
390
391
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

392
393
394
395
396
397
398
399
400
401
402
403
404
			tokenBatch.Clear()
			embedBatch.Clear()
		}
	}
}

// TODO (jmorganca): processBatch should be simplified, removing:
// * sampling
// * stop token checking
// * metrics
// these should instead be handled by the handlers
// it should only be responsible for accepting tokens or embeddings and
// processing batches as fast as possible
405
func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch) error {
406
407
408
409
410
411
412
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

	var batch *llama.Batch
413
	var numOutputs int
414
415
416
417
418
419
420
421
422
423
424

	seqIdx := s.nextSeq - 1
	for range s.seqs {
		seqIdx = (seqIdx + 1) % len(s.seqs)
		seq := s.seqs[seqIdx]

		if seq == nil {
			continue
		}

		// if past the num predict limit
425
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
426
			s.removeSequence(seqIdx, llm.DoneReasonLength)
427
428
429
430
			continue
		}

		for i, input := range seq.inputs {
431
432
			if len(seq.cache.Inputs)+len(seq.pendingInputs)+1 > s.cache.numCtx {
				if len(seq.pendingInputs) == 0 {
433
434
435
436
437
					if !seq.shift {
						s.removeSequence(seqIdx, llm.DoneReasonLength)
						break
					}

438
439
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
440
441
442
443
444
445
446
447
448
						var reprocess *ErrReprocessInputs
						if errors.As(err, &reprocess) {
							// Prepend these inputs to the sequence's inputs queue for reprocessing
							seq.inputs = append(reprocess.Inputs, seq.inputs...)
							// Continue processing as normal
							continue
						} else {
							return err
						}
449
					}
450
451
452
453
454
				} else {
					break
				}
			}

455
456
457
458
459
460
461
462
463
464
465
466
			embedding := input.embed != nil

			// If we don't currently have a batch, use one of the correct type and
			// fill it up as much as possible across all sequences. If we encounter an
			// input of the opppsite type, stop for that sequence but then pick up from
			// there for the next batch, ensuring that we alternate types
			if batch == nil {
				if !embedding {
					batch = tokenBatch
				} else {
					batch = embedBatch
				}
467
			} else if embedding != batch.IsEmbedding() {
468
469
470
471
				s.nextSeq = seqIdx
				break
			}

472
			if i >= batch.Size() {
473
474
475
				break
			}

476
477
478
479
480
481
			output := i+1 == len(seq.inputs)
			batch.Add(input.token, input.embed, len(seq.cache.Inputs)+len(seq.pendingInputs), output, seq.cache.Id)
			if output {
				numOutputs++
			}

482
			seq.pendingInputs = append(seq.pendingInputs, input)
483
484
			seq.iBatch = batch.NumTokens() - 1
		}
485
486

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
487
488
489
	}

	if batch == nil || batch.NumTokens() == 0 {
490
		return nil
491
492
	}

Michael Yang's avatar
Michael Yang committed
493
494
	t := time.Now()
	if err := s.lc.Decode(batch); err != nil {
495
		return fmt.Errorf("failed to decode batch: %w", err)
496
497
	}

498
499
500
501
	if numOutputs > 0 {
		s.lc.Synchronize()
	}

502
503
504
505
506
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

507
508
509
510
511
512
		// After calling Decode, pending inputs are now in the cache
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
			seq.pendingInputs = []input{}
		}

513
514
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
515
			seq.processingDuration += time.Since(t)
516
517
518
			continue
		}

Michael Yang's avatar
Michael Yang committed
519
520
521
522
523
		seq.numDecoded++
		if seq.numDecoded > 1 {
			seq.generationDuration += time.Since(t)
		} else {
			seq.processingDuration += time.Since(t)
524
525
526
527
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
528
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
529
530
531
532
533
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
534
			s.removeSequence(i, llm.DoneReasonStop)
535
536
537
538
			continue
		}

		// sample a token
539
540
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
541
542
543
544
545
546
547
548
549
550
		piece := s.model.TokenToPiece(token)

		seq.numPredicted++

		// if it's an end of sequence token, break
		if s.model.TokenIsEog(token) {
			// TODO (jmorganca): we should send this back
			// as it's important for the /api/generate context
			// seq.responses <- piece

551
			s.removeSequence(i, llm.DoneReasonStop)
552
553
554
			continue
		}

555
556
557
558
559
560
561
562
563
		// Calculate logprobs if requested (after EOS check to avoid logprobs for EOS tokens)
		if seq.logprobs {
			logits := s.lc.GetLogitsIth(seq.iBatch)
			if logits != nil {
				logprobs := calculateLogprobsLlama(logits, token, seq.topLogprobs, s.model)
				seq.pendingLogprobs = append(seq.pendingLogprobs, logprobs...)
			}
		}

564
565
566
567
568
		seq.inputs = []input{{token: token}}

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

Jesse Gross's avatar
Jesse Gross committed
569
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
570
571
572
573
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
574
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
575
576
			newLen := len(seq.pendingResponses)

577
578
579
580
581
582
583
584
585
586
587
			// Truncate logprobs to match the truncated responses
			if seq.logprobs {
				origLogprobsLen := len(seq.pendingLogprobs)
				numTokensRemoved := origLen - newLen
				newLogprobsLen := origLogprobsLen - numTokensRemoved
				if newLogprobsLen < 0 {
					newLogprobsLen = 0
				}
				seq.pendingLogprobs = seq.pendingLogprobs[:newLogprobsLen]
			}

588
589
590
591
592
593
594
595
596
597
598
599
600
			// 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]
601

602
			s.removeSequence(i, llm.DoneReasonStop)
603
604
605
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
606
		if common.ContainsStopSuffix(sequence, seq.stop) {
607
608
609
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
610
		if common.IncompleteUnicode(sequence) {
611
612
613
614
			continue
		}

		if !flushPending(seq) {
615
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
616
617
		}
	}
618
619

	return nil
620
621
622
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
623
	var req llm.CompletionRequest
624
625
626
627
628
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

629
630
631
632
633
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

634
635
636
637
638
639
640
641
642
643
	// 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
	}

644
645
646
647
648
649
650
651
652
653
654
655
656
657
	// Extract options from the CompletionRequest
	samplingParams := llama.SamplingParams{
		TopK:           req.Options.TopK,
		TopP:           req.Options.TopP,
		MinP:           req.Options.MinP,
		TypicalP:       req.Options.TypicalP,
		Temp:           req.Options.Temperature,
		RepeatLastN:    req.Options.RepeatLastN,
		PenaltyRepeat:  req.Options.RepeatPenalty,
		PenaltyFreq:    req.Options.FrequencyPenalty,
		PenaltyPresent: req.Options.PresencePenalty,
		Seed:           uint32(req.Options.Seed),
		Grammar:        req.Grammar,
	}
658
659

	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
660
661
662
		numPredict:     req.Options.NumPredict,
		stop:           req.Options.Stop,
		numKeep:        req.Options.NumKeep,
663
664
		samplingParams: &samplingParams,
		embedding:      false,
665
666
		shift:          req.Shift,
		truncate:       req.Truncate,
667
668
		logprobs:       req.Logprobs,
		topLogprobs:    req.TopLogprobs,
669
670
	})
	if err != nil {
671
672
673
674
		if errors.Is(err, errorInputTooLong) {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
675
676
677
678
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

679
	// Ensure there is a place to put the sequence, released when removed from s.seqs
680
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
681
682
683
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
684
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
685
		}
686
687
688
		return
	}

689
	s.mu.Lock()
690
	found := false
691
692
	for i, sq := range s.seqs {
		if sq == nil {
693
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
694
695
			if err != nil {
				s.mu.Unlock()
696
				s.seqsSem.Release(1)
697
698
699
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
700

701
702
			s.seqs[i] = seq
			s.cond.Signal()
703
			found = true
704
705
706
707
708
			break
		}
	}
	s.mu.Unlock()

709
	if !found {
710
		s.seqsSem.Release(1)
711
712
713
714
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

715
716
717
718
719
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
720
		case resp, ok := <-seq.responses:
721
			if ok {
722
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
723
724
					Content:  resp.content,
					Logprobs: resp.logprobs,
725
726
727
728
729
730
731
732
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
733
734
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
735
					DoneReason:         seq.doneReason,
736
					PromptEvalCount:    seq.numPromptInputs,
Michael Yang's avatar
Michael Yang committed
737
					PromptEvalDuration: seq.processingDuration,
738
					EvalCount:          seq.numDecoded,
Michael Yang's avatar
Michael Yang committed
739
					EvalDuration:       seq.generationDuration,
740
741
742
743
744
745
746
747
748
749
750
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
				}

				return
			}
		}
	}
}

func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
751
	var req llm.EmbeddingRequest
752
753
754
755
756
757
758
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, fmt.Sprintf("bad request: %s", err), http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/json")

759
760
	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{
		embedding: true,
761
		truncate:  false,
762
	})
763
	if err != nil {
764
765
766
767
		if errors.Is(err, errorInputTooLong) {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
768
769
770
771
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

772
	// Ensure there is a place to put the sequence, released when removed from s.seqs
773
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
774
775
776
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
777
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
778
		}
779
780
781
		return
	}

782
	s.mu.Lock()
783
	found := false
784
785
	for i, sq := range s.seqs {
		if sq == nil {
786
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
787
788
			if err != nil {
				s.mu.Unlock()
789
				s.seqsSem.Release(1)
790
791
792
793
794
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
			s.seqs[i] = seq
			s.cond.Signal()
795
			found = true
796
797
798
799
800
			break
		}
	}
	s.mu.Unlock()

801
	if !found {
802
		s.seqsSem.Release(1)
803
804
805
806
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

807
808
	embedding := <-seq.embedding

809
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
810
811
		Embedding:       embedding,
		PromptEvalCount: seq.numPromptInputs,
812
813
814
815
816
817
818
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
819
820
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
821
822
823
824
825
826
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

Jesse Gross's avatar
Jesse Gross committed
827
828
// loadModel allocates memory based on the given parameters and loads the weights. The
// memory allocated is worst case for text models but not for vision.
829
830
831
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
Jesse Gross's avatar
Jesse Gross committed
832
	lpath []string,
833
834
	ppath string,
	kvSize int,
835
	kvCacheType string,
836
	flashAttention ml.FlashAttentionType,
837
838
839
	threads int,
	multiUserCache bool,
) {
840
841
842
843
844
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
845

846
	ctxParams := llama.NewContextParams(kvSize, s.batchSize, s.parallel, threads, flashAttention, kvCacheType)
847
848
849
850
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
851

Jesse Gross's avatar
Jesse Gross committed
852
853
854
855
	for _, path := range lpath {
		err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
		if err != nil {
			panic(err)
856
857
858
859
		}
	}

	if ppath != "" {
860
		var err error
861
		s.image, err = NewImageContext(s.lc, ppath)
862
863
864
		if err != nil {
			panic(err)
		}
865
866
	}

867
868
869
870
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
871

872
	s.status = llm.ServerStatusReady
873
874
875
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
// load is the handler called by the Ollama server to process different
// load operations
func (s *Server) load(w http.ResponseWriter, r *http.Request) {
	s.loadMu.Lock()
	defer s.loadMu.Unlock()

	w.Header().Set("Content-Type", "application/json")

	if s.status != llm.ServerStatusLaunched {
		http.Error(w, "model already loaded", http.StatusInternalServerError)
		return
	}

	var req llm.LoadRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "bad request", http.StatusBadRequest)
		return
	}

	slog.Info("load", "request", req)

	switch req.Operation {
	// LoadOperationFit and LoadOperationAlloc have no meaning here - just return a successful response

	case llm.LoadOperationCommit:
		s.batchSize = req.BatchSize
		s.parallel = req.Parallel
		s.seqs = make([]*Sequence, s.parallel)
		s.seqsSem = semaphore.NewWeighted(int64(s.parallel))

		numGPU := 0
907
908
909
910
911
912
913
914
		var tensorSplit []float32
		var llamaIDs []uint64

		gpuIDs := llama.EnumerateGPUs()
		sort.Sort(req.GPULayers)
		for _, layers := range req.GPULayers {
			for i := range gpuIDs {
				if gpuIDs[i].DeviceID == layers.DeviceID {
Jesse Gross's avatar
Jesse Gross committed
915
					numGPU += len(layers.Layers)
916
917
					tensorSplit = append(tensorSplit, float32(len(layers.Layers)))
					llamaIDs = append(llamaIDs, gpuIDs[i].LlamaID)
Jesse Gross's avatar
Jesse Gross committed
918
919
920
921
922
				}
			}
		}

		params := llama.ModelParams{
923
			Devices:      llamaIDs,
Jesse Gross's avatar
Jesse Gross committed
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
			NumGpuLayers: numGPU,
			MainGpu:      req.MainGPU,
			UseMmap:      req.UseMmap && len(req.LoraPath) == 0,
			TensorSplit:  tensorSplit,
			Progress: func(progress float32) {
				s.progress = progress
			},
		}

		s.status = llm.ServerStatusLoadingModel
		go s.loadModel(params, s.modelPath, req.LoraPath, req.ProjectorPath, req.KvSize, req.KvCacheType, req.FlashAttention, req.NumThreads, req.MultiUserCache)

	case llm.LoadOperationClose:
		// No-op for us
		if err := json.NewEncoder(w).Encode(&llm.LoadResponse{}); err != nil {
			http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		}
		return
	}

	resp := llm.LoadResponse{Success: true}
	if err := json.NewEncoder(w).Encode(&resp); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		return
	}
}

951
952
953
954
func Execute(args []string) error {
	fs := flag.NewFlagSet("runner", flag.ExitOnError)
	mpath := fs.String("model", "", "Path to model binary file")
	port := fs.Int("port", 8080, "Port to expose the server on")
955
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
956

957
958
959
960
961
962
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
963
	}
964
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
965
	slog.Info("starting go runner")
966
967

	llama.BackendInit()
968
969

	server := &Server{
Jesse Gross's avatar
Jesse Gross committed
970
971
		modelPath: *mpath,
		status:    llm.ServerStatusLaunched,
972
973
974
975
976
977
978
	}

	server.ready.Add(1)

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

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

981
982
983
984
985
986
	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)
987
		return err
988
989
990
991
	}
	defer listener.Close()

	mux := http.NewServeMux()
Jesse Gross's avatar
Jesse Gross committed
992
	mux.HandleFunc("POST /load", server.load)
993
994
995
996
997
998
999
1000
1001
1002
1003
	mux.HandleFunc("/embedding", server.embeddings)
	mux.HandleFunc("/completion", server.completion)
	mux.HandleFunc("/health", server.health)

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

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

1007
	return nil
1008
}