runner.go 22.9 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
166
167
168
169
170
	}, 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
171
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
172
	var inputs []input
173
174
175
176
177
178
179
180
181
182
	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}
	}
183
184
185

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

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

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

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

216
217
218
219
220
221
222
223
224
225
			for _, e := range embed {
				inputs = append(inputs, input{embed: e})
			}
		}
	}

	return inputs, nil
}

type Server struct {
Jesse Gross's avatar
Jesse Gross committed
226
227
228
229
230
231
	// 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

232
233
234
235
236
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
237
238
	model *llama.Model

239
	// image model context for multi-modal models
240
	image *ImageContext
241

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

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

255
256
257
258
259
260
261
262
263
	// 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
264

265
	// the list of simultaneous sequences being evaluated
266
267
	seqs []*Sequence

268
269
270
271
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
	// 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 {
289
290
291
292
293
294
295
296
297
298
299
	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]
300
301
	}

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

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
312
313
}

314
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
315
316
317
318
319
320
321
322
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
323
	s.seqsSem.Release(1)
324
325
326
327
328
}

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

329
	// Logically these batches are used only within the context of processBatch
330
	// but it is better for performance to allocate them once here
Jesse Gross's avatar
Jesse Gross committed
331
332
333
334
	tokenBatch, err := llama.NewBatch(s.batchSize, len(s.seqs), 0)
	if err != nil {
		panic(err)
	}
335
336
	defer tokenBatch.Free()

337
338
339
	var embedBatch *llama.Batch
	embedBatchSize := s.image.BatchSize(s.batchSize)
	if embedBatchSize != 0 {
Jesse Gross's avatar
Jesse Gross committed
340
341
342
343
		embedBatch, err = llama.NewBatch(embedBatchSize, len(s.seqs), s.image.EmbedSize(s.lc))
		if err != nil {
			panic(err)
		}
344
345
346
347
		defer embedBatch.Free()
	} else {
		embedBatch = &llama.Batch{}
	}
348
349
350
351
352
353

	for {
		select {
		case <-ctx.Done():
			return
		default:
354
355
356
357
358
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

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

	var batch *llama.Batch

	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
391
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
392
			s.removeSequence(seqIdx, llm.DoneReasonLength)
393
394
395
396
			continue
		}

		for i, input := range seq.inputs {
397
398
			if len(seq.cache.Inputs)+len(seq.pendingInputs)+1 > s.cache.numCtx {
				if len(seq.pendingInputs) == 0 {
399
400
401
402
403
					if !seq.shift {
						s.removeSequence(seqIdx, llm.DoneReasonLength)
						break
					}

404
405
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
406
407
408
409
410
411
412
413
414
						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
						}
415
					}
416
417
418
419
420
				} else {
					break
				}
			}

421
422
423
424
425
426
427
428
429
430
431
432
			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
				}
433
			} else if embedding != batch.IsEmbedding() {
434
435
436
437
				s.nextSeq = seqIdx
				break
			}

438
			if i >= batch.Size() {
439
440
441
				break
			}

442
443
			batch.Add(input.token, input.embed, len(seq.cache.Inputs)+len(seq.pendingInputs), i+1 == len(seq.inputs), seq.cache.Id)
			seq.pendingInputs = append(seq.pendingInputs, input)
444
445
			seq.iBatch = batch.NumTokens() - 1
		}
446
447

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
448
449
450
	}

	if batch == nil || batch.NumTokens() == 0 {
451
		return nil
452
453
	}

Michael Yang's avatar
Michael Yang committed
454
455
	t := time.Now()
	if err := s.lc.Decode(batch); err != nil {
456
		return fmt.Errorf("failed to decode batch: %w", err)
457
458
459
460
461
462
463
	}

	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

464
465
466
467
468
469
		// 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{}
		}

470
471
472
473
474
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
			continue
		}

Michael Yang's avatar
Michael Yang committed
475
476
477
478
479
480
		s.lc.Synchronize()
		seq.numDecoded++
		if seq.numDecoded > 1 {
			seq.generationDuration += time.Since(t)
		} else {
			seq.processingDuration += time.Since(t)
481
482
483
484
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
485
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
486
487
488
489
490
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
491
			s.removeSequence(i, llm.DoneReasonStop)
492
493
494
495
			continue
		}

		// sample a token
496
497
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
498
499
500
501
502
503
504
505
506
507
		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

508
			s.removeSequence(i, llm.DoneReasonStop)
509
510
511
512
513
514
515
516
			continue
		}

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

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

Jesse Gross's avatar
Jesse Gross committed
517
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
518
519
520
521
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

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

539
			s.removeSequence(i, llm.DoneReasonStop)
540
541
542
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
543
		if common.ContainsStopSuffix(sequence, seq.stop) {
544
545
546
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
547
		if common.IncompleteUnicode(sequence) {
548
549
550
551
			continue
		}

		if !flushPending(seq) {
552
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
553
554
		}
	}
555
556

	return nil
557
558
559
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
560
	var req llm.CompletionRequest
561
562
563
564
565
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

566
567
568
569
570
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

571
572
573
574
575
576
577
578
579
580
	// 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
	}

581
582
583
584
585
586
587
588
589
590
591
592
593
594
	// 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,
	}
595
596

	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
597
598
599
		numPredict:     req.Options.NumPredict,
		stop:           req.Options.Stop,
		numKeep:        req.Options.NumKeep,
600
601
		samplingParams: &samplingParams,
		embedding:      false,
602
603
		shift:          req.Shift,
		truncate:       req.Truncate,
604
605
	})
	if err != nil {
606
607
608
609
		if errors.Is(err, errorInputTooLong) {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
610
611
612
613
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

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

624
	s.mu.Lock()
625
	found := false
626
627
	for i, sq := range s.seqs {
		if sq == nil {
628
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
629
630
			if err != nil {
				s.mu.Unlock()
631
				s.seqsSem.Release(1)
632
633
634
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
635

636
637
			s.seqs[i] = seq
			s.cond.Signal()
638
			found = true
639
640
641
642
643
			break
		}
	}
	s.mu.Unlock()

644
	if !found {
645
		s.seqsSem.Release(1)
646
647
648
649
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

650
651
652
653
654
655
656
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
657
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
658
659
660
661
662
663
664
665
666
					Content: content,
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
667
668
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
669
					DoneReason:         seq.doneReason,
670
					PromptEvalCount:    seq.numPromptInputs,
Michael Yang's avatar
Michael Yang committed
671
					PromptEvalDuration: seq.processingDuration,
672
					EvalCount:          seq.numDecoded,
Michael Yang's avatar
Michael Yang committed
673
					EvalDuration:       seq.generationDuration,
674
675
676
677
678
679
680
681
682
683
684
				}); 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) {
685
	var req llm.EmbeddingRequest
686
687
688
689
690
691
692
693
694
695
696
697
698
	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")

	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{embedding: true})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

699
	// Ensure there is a place to put the sequence, released when removed from s.seqs
700
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
701
702
703
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
704
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
705
		}
706
707
708
		return
	}

709
	s.mu.Lock()
710
	found := false
711
712
	for i, sq := range s.seqs {
		if sq == nil {
713
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
714
715
			if err != nil {
				s.mu.Unlock()
716
				s.seqsSem.Release(1)
717
718
719
720
721
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
			s.seqs[i] = seq
			s.cond.Signal()
722
			found = true
723
724
725
726
727
			break
		}
	}
	s.mu.Unlock()

728
	if !found {
729
		s.seqsSem.Release(1)
730
731
732
733
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

734
735
	embedding := <-seq.embedding

736
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
737
738
739
740
741
742
743
744
		Embedding: embedding,
	}); 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")
745
746
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
747
748
749
750
751
752
		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
753
754
// 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.
755
756
757
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
Jesse Gross's avatar
Jesse Gross committed
758
	lpath []string,
759
760
	ppath string,
	kvSize int,
761
	kvCacheType string,
762
763
764
765
	flashAttention bool,
	threads int,
	multiUserCache bool,
) {
766
767
768
769
770
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
771

772
	ctxParams := llama.NewContextParams(kvSize, s.batchSize*s.parallel, s.parallel, threads, flashAttention, kvCacheType)
773
774
775
776
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
777

Jesse Gross's avatar
Jesse Gross committed
778
779
780
781
	for _, path := range lpath {
		err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
		if err != nil {
			panic(err)
782
783
784
785
		}
	}

	if ppath != "" {
786
		var err error
787
		s.image, err = NewImageContext(s.lc, ppath)
788
789
790
		if err != nil {
			panic(err)
		}
791
792
	}

793
794
795
796
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
797

798
	s.status = llm.ServerStatusReady
799
800
801
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
// 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 {
837
				if gpuIDs[i] == layers.DeviceID {
Jesse Gross's avatar
Jesse Gross committed
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
					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
	}
}

872
873
874
875
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")
876
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
877

878
879
880
881
882
883
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
884
	}
885
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
886
	slog.Info("starting go runner")
887
888

	llama.BackendInit()
889
890

	server := &Server{
Jesse Gross's avatar
Jesse Gross committed
891
892
		modelPath: *mpath,
		status:    llm.ServerStatusLaunched,
893
894
895
896
897
898
899
	}

	server.ready.Add(1)

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

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

902
903
904
905
906
907
	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)
908
		return err
909
910
911
912
	}
	defer listener.Close()

	mux := http.NewServeMux()
Jesse Gross's avatar
Jesse Gross committed
913
	mux.HandleFunc("POST /load", server.load)
914
915
916
917
918
919
920
921
922
923
924
	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)
925
		return err
926
927
	}

928
	return nil
929
}