runner.go 23.5 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
15
16
17
18

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

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

23
	"github.com/ollama/ollama/api"
24
	"github.com/ollama/ollama/envconfig"
25
	"github.com/ollama/ollama/llama"
26
	"github.com/ollama/ollama/llm"
27
	"github.com/ollama/ollama/logutil"
Jesse Gross's avatar
Jesse Gross committed
28
	"github.com/ollama/ollama/runner/common"
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
)

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

50
51
52
	// inputs that have been added to a batch but not yet submitted to Decode
	pendingInputs []input

53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
	// 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

	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

82
83
84
	// shift if context window is exceeded
	shift bool

85
	doneReason llm.DoneReason
86
87

	// Metrics
Michael Yang's avatar
Michael Yang committed
88
89
90
91
	processingDuration time.Duration
	generationDuration time.Duration
	numDecoded         int
	numPromptInputs    int
92
93
94
95
96
97
98
99
}

type NewSequenceParams struct {
	numPredict     int
	stop           []string
	numKeep        int
	samplingParams *llama.SamplingParams
	embedding      bool
100
101
	shift          bool
	truncate       bool
102
103
}

104
105
var errorInputTooLong = errors.New("the input length exceeds the context length")

106
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
107
108
109
110
111
112
113
114
115
116
117
118
119
	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)
	}

120
121
	if s.model.AddBOSToken() {
		params.numKeep += 1
122
123
	}

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

127
	if len(inputs) > s.cache.numCtx {
128
		discard := len(inputs) - s.cache.numCtx
129
130
131
132
		if !params.truncate {
			return nil, errorInputTooLong
		}

133
		newInputs := inputs[:params.numKeep]
134
135
136
		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))
137
		inputs = newInputs
138
139
140
141
	}

	var sc *llama.SamplingContext
	if params.samplingParams != nil {
Jesse Gross's avatar
Jesse Gross committed
142
143
144
145
		sc, err = llama.NewSamplingContext(s.model, *params.samplingParams)
		if err != nil {
			return nil, err
		}
146
147
		for _, input := range inputs {
			if input.embed == nil {
148
				sc.Accept(input.token, false)
149
150
151
152
153
			}
		}
	}

	return &Sequence{
Michael Yang's avatar
Michael Yang committed
154
155
156
157
158
159
160
161
162
163
164
		inputs:           inputs,
		numPromptInputs:  len(inputs),
		numPredict:       params.numPredict,
		pendingResponses: make([]string, 0),
		responses:        make(chan string, 100),
		quit:             make(chan bool, 1),
		embedding:        make(chan []float32, 1),
		samplingCtx:      sc,
		embeddingOnly:    params.embedding,
		stop:             params.stop,
		numKeep:          params.numKeep,
165
		shift:            params.shift,
166
167
168
169
170
171
	}, nil
}

// 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
172
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
173
	var inputs []input
174
175
176
177
178
179
180
181
182
183
	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}
	}
184
185
186

	for i, part := range parts {
		// text - tokenize
187
188
189
190
		tokens, err := s.lc.Model().Tokenize(part, i == 0, true)
		if err != nil {
			return nil, err
		}
191

192
193
		for _, t := range tokens {
			inputs = append(inputs, input{token: t})
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
		}

		// 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)
			}

212
			chunks, err := s.image.MultimodalTokenize(s.lc, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
213
214
215
216
			if err != nil {
				return nil, err
			}

217
218
219
220
221
222
223
224
			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})
					}
				}
225
226
227
228
229
230
231
232
			}
		}
	}

	return inputs, nil
}

type Server struct {
Jesse Gross's avatar
Jesse Gross committed
233
234
235
236
237
238
	// 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

239
240
241
242
243
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
244
245
	model *llama.Model

246
	// image model context for multi-modal models
247
	image *ImageContext
248

249
	// status for external health reporting - loading, ready to serve, etc.
250
	status llm.ServerStatus
251
252
253
254
255
256
257
258

	// 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)
259
	// TODO (jmorganca): make this n_batch
260
261
	batchSize int

262
263
264
265
266
267
268
269
270
	// 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
271

272
	// the list of simultaneous sequences being evaluated
273
274
	seqs []*Sequence

275
276
277
278
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
	// 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 {
296
297
298
299
300
301
302
303
304
305
306
	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]
307
308
	}

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

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
319
320
}

321
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
322
323
324
325
326
327
328
329
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
330
	s.seqsSem.Release(1)
331
332
333
334
335
}

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

336
	// Logically these batches are used only within the context of processBatch
337
	// but it is better for performance to allocate them once here
Jesse Gross's avatar
Jesse Gross committed
338
339
340
341
	tokenBatch, err := llama.NewBatch(s.batchSize, len(s.seqs), 0)
	if err != nil {
		panic(err)
	}
342
343
	defer tokenBatch.Free()

344
345
346
	var embedBatch *llama.Batch
	embedBatchSize := s.image.BatchSize(s.batchSize)
	if embedBatchSize != 0 {
Jesse Gross's avatar
Jesse Gross committed
347
348
349
350
		embedBatch, err = llama.NewBatch(embedBatchSize, len(s.seqs), s.image.EmbedSize(s.lc))
		if err != nil {
			panic(err)
		}
351
352
353
354
		defer embedBatch.Free()
	} else {
		embedBatch = &llama.Batch{}
	}
355
356
357
358
359
360

	for {
		select {
		case <-ctx.Done():
			return
		default:
361
362
363
364
365
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

366
367
368
369
370
371
372
373
374
375
376
377
378
			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
379
func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch) error {
380
381
382
383
384
385
386
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

	var batch *llama.Batch
387
	var numOutputs int
388
389
390
391
392
393
394
395
396
397
398

	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
399
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
400
			s.removeSequence(seqIdx, llm.DoneReasonLength)
401
402
403
404
			continue
		}

		for i, input := range seq.inputs {
405
406
			if len(seq.cache.Inputs)+len(seq.pendingInputs)+1 > s.cache.numCtx {
				if len(seq.pendingInputs) == 0 {
407
408
409
410
411
					if !seq.shift {
						s.removeSequence(seqIdx, llm.DoneReasonLength)
						break
					}

412
413
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
414
415
416
417
418
419
420
421
422
						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
						}
423
					}
424
425
426
427
428
				} else {
					break
				}
			}

429
430
431
432
433
434
435
436
437
438
439
440
			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
				}
441
			} else if embedding != batch.IsEmbedding() {
442
443
444
445
				s.nextSeq = seqIdx
				break
			}

446
			if i >= batch.Size() {
447
448
449
				break
			}

450
451
452
453
454
455
			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++
			}

456
			seq.pendingInputs = append(seq.pendingInputs, input)
457
458
			seq.iBatch = batch.NumTokens() - 1
		}
459
460

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
461
462
463
	}

	if batch == nil || batch.NumTokens() == 0 {
464
		return nil
465
466
	}

Michael Yang's avatar
Michael Yang committed
467
468
	t := time.Now()
	if err := s.lc.Decode(batch); err != nil {
469
		return fmt.Errorf("failed to decode batch: %w", err)
470
471
	}

472
473
474
475
	if numOutputs > 0 {
		s.lc.Synchronize()
	}

476
477
478
479
480
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

481
482
483
484
485
486
		// 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{}
		}

487
488
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
489
			seq.processingDuration += time.Since(t)
490
491
492
			continue
		}

Michael Yang's avatar
Michael Yang committed
493
494
495
496
497
		seq.numDecoded++
		if seq.numDecoded > 1 {
			seq.generationDuration += time.Since(t)
		} else {
			seq.processingDuration += time.Since(t)
498
499
500
501
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
502
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
503
504
505
506
507
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
508
			s.removeSequence(i, llm.DoneReasonStop)
509
510
511
512
			continue
		}

		// sample a token
513
514
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
515
516
517
518
519
520
521
522
523
524
		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

525
			s.removeSequence(i, llm.DoneReasonStop)
526
527
528
529
530
531
532
533
			continue
		}

		seq.inputs = []input{{token: token}}

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

Jesse Gross's avatar
Jesse Gross committed
534
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
535
536
537
538
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
539
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
			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]
555

556
			s.removeSequence(i, llm.DoneReasonStop)
557
558
559
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
560
		if common.ContainsStopSuffix(sequence, seq.stop) {
561
562
563
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
564
		if common.IncompleteUnicode(sequence) {
565
566
567
568
			continue
		}

		if !flushPending(seq) {
569
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
570
571
		}
	}
572
573

	return nil
574
575
576
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
577
	var req llm.CompletionRequest
578
579
580
581
582
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

583
584
585
586
587
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

588
589
590
591
592
593
594
595
596
597
	// 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
	}

598
599
600
601
602
603
604
605
606
607
608
609
610
611
	// 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,
	}
612
613

	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
614
615
616
		numPredict:     req.Options.NumPredict,
		stop:           req.Options.Stop,
		numKeep:        req.Options.NumKeep,
617
618
		samplingParams: &samplingParams,
		embedding:      false,
619
620
		shift:          req.Shift,
		truncate:       req.Truncate,
621
622
	})
	if err != nil {
623
624
625
626
		if errors.Is(err, errorInputTooLong) {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
627
628
629
630
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

631
	// Ensure there is a place to put the sequence, released when removed from s.seqs
632
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
633
634
635
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
636
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
637
		}
638
639
640
		return
	}

641
	s.mu.Lock()
642
	found := false
643
644
	for i, sq := range s.seqs {
		if sq == nil {
645
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
646
647
			if err != nil {
				s.mu.Unlock()
648
				s.seqsSem.Release(1)
649
650
651
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
652

653
654
			s.seqs[i] = seq
			s.cond.Signal()
655
			found = true
656
657
658
659
660
			break
		}
	}
	s.mu.Unlock()

661
	if !found {
662
		s.seqsSem.Release(1)
663
664
665
666
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

667
668
669
670
671
672
673
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
674
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
675
676
677
678
679
680
681
682
683
					Content: content,
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
684
685
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
686
					DoneReason:         seq.doneReason,
687
					PromptEvalCount:    seq.numPromptInputs,
Michael Yang's avatar
Michael Yang committed
688
					PromptEvalDuration: seq.processingDuration,
689
					EvalCount:          seq.numDecoded,
Michael Yang's avatar
Michael Yang committed
690
					EvalDuration:       seq.generationDuration,
691
692
693
694
695
696
697
698
699
700
701
				}); 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) {
702
	var req llm.EmbeddingRequest
703
704
705
706
707
708
709
	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")

710
711
	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{
		embedding: true,
712
713
714
715
716

		// TODO (jmorganca): this should be provided by the server via the
		// request options and truncated here in the runner, instead of relying on
		// the server's truncate logic
		truncate: true,
717
	})
718
719
720
721
722
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

723
	// Ensure there is a place to put the sequence, released when removed from s.seqs
724
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
725
726
727
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
728
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
729
		}
730
731
732
		return
	}

733
	s.mu.Lock()
734
	found := false
735
736
	for i, sq := range s.seqs {
		if sq == nil {
737
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
738
739
			if err != nil {
				s.mu.Unlock()
740
				s.seqsSem.Release(1)
741
742
743
744
745
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
			s.seqs[i] = seq
			s.cond.Signal()
746
			found = true
747
748
749
750
751
			break
		}
	}
	s.mu.Unlock()

752
	if !found {
753
		s.seqsSem.Release(1)
754
755
756
757
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

758
759
	embedding := <-seq.embedding

760
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
761
		Embedding: embedding,
762
763
764
765
766
767
768
	}); 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")
769
770
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
771
772
773
774
775
776
		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
777
778
// 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.
779
780
781
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
Jesse Gross's avatar
Jesse Gross committed
782
	lpath []string,
783
784
	ppath string,
	kvSize int,
785
	kvCacheType string,
786
787
788
789
	flashAttention bool,
	threads int,
	multiUserCache bool,
) {
790
791
792
793
794
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
795

796
	ctxParams := llama.NewContextParams(kvSize, s.batchSize*s.parallel, s.parallel, threads, flashAttention, kvCacheType)
797
798
799
800
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
801

Jesse Gross's avatar
Jesse Gross committed
802
803
804
805
	for _, path := range lpath {
		err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
		if err != nil {
			panic(err)
806
807
808
809
		}
	}

	if ppath != "" {
810
		var err error
811
		s.image, err = NewImageContext(s.lc, ppath)
812
813
814
		if err != nil {
			panic(err)
		}
815
816
	}

817
818
819
820
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
821

822
	s.status = llm.ServerStatusReady
823
824
825
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
// 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))

		gpuIDs := llama.EnumerateGPUs()
		tensorSplit := make([]float32, len(gpuIDs))
		numGPU := 0
		for i := range gpuIDs {
			for _, layers := range req.GPULayers {
861
				if gpuIDs[i] == layers.DeviceID {
Jesse Gross's avatar
Jesse Gross committed
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
					tensorSplit[i] = float32(len(layers.Layers))
					numGPU += len(layers.Layers)
				}
			}
		}

		params := llama.ModelParams{
			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
	}
}

896
897
898
899
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")
900
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
901

902
903
904
905
906
907
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
908
	}
909
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
910
	slog.Info("starting go runner")
911
912

	llama.BackendInit()
913
914

	server := &Server{
Jesse Gross's avatar
Jesse Gross committed
915
916
		modelPath: *mpath,
		status:    llm.ServerStatusLaunched,
917
918
919
920
921
922
923
	}

	server.ready.Add(1)

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

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

926
927
928
929
930
931
	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)
932
		return err
933
934
935
936
	}
	defer listener.Close()

	mux := http.NewServeMux()
Jesse Gross's avatar
Jesse Gross committed
937
	mux.HandleFunc("POST /load", server.load)
938
939
940
941
942
943
944
945
946
947
948
	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)
949
		return err
950
951
	}

952
	return nil
953
}