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

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"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"
Jesse Gross's avatar
Jesse Gross committed
29
	"github.com/ollama/ollama/runner/common"
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
)

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

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

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

83
	doneReason llm.DoneReason
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

	// Metrics
	startProcessingTime time.Time
	startGenerationTime time.Time
	numDecoded          int
	numPromptInputs     int
}

type NewSequenceParams struct {
	numPredict     int
	stop           []string
	numKeep        int
	samplingParams *llama.SamplingParams
	embedding      bool
}

100
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
	s.ready.Wait()

	startTime := time.Now()

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

116
117
	if s.model.AddBOSToken() {
		params.numKeep += 1
118
119
	}

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

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

	var sc *llama.SamplingContext
	if params.samplingParams != nil {
Jesse Gross's avatar
Jesse Gross committed
134
135
136
137
		sc, err = llama.NewSamplingContext(s.model, *params.samplingParams)
		if err != nil {
			return nil, err
		}
138
139
		for _, input := range inputs {
			if input.embed == nil {
140
				sc.Accept(input.token, false)
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
			}
		}
	}

	return &Sequence{
		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),
		samplingCtx:         sc,
		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
// generating image embeddings for each image
164
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
165
	var inputs []input
166
167
168
169
170
171
172
173
174
175
	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}
	}
176
177
178

	for i, part := range parts {
		// text - tokenize
179
180
181
182
		tokens, err := s.lc.Model().Tokenize(part, i == 0, true)
		if err != nil {
			return nil, err
		}
183

184
185
		for _, t := range tokens {
			inputs = append(inputs, input{token: t})
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
		}

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

204
			embed, err := s.image.NewEmbed(s.lc, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
205
206
207
208
			if err != nil {
				return nil, err
			}

209
210
211
212
213
214
215
216
217
218
			for _, e := range embed {
				inputs = append(inputs, input{embed: e})
			}
		}
	}

	return inputs, nil
}

type Server struct {
219
220
221
222
223
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
224
225
	model *llama.Model

226
	// image model context for multi-modal models
227
	image *ImageContext
228

229
	// status for external health reporting - loading, ready to serve, etc.
230
	status llm.ServerStatus
231
232
233
234
235
236
237
238

	// 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)
239
	// TODO (jmorganca): make this n_batch
240
241
	batchSize int

242
243
244
245
246
247
248
249
250
	// 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
251

252
	// the list of simultaneous sequences being evaluated
253
254
	seqs []*Sequence

255
256
257
258
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
	// 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 {
276
277
278
279
280
281
282
283
284
285
286
	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]
287
288
	}

289
290
291
292
293
294
295
296
297
298
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
299
300
}

301
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
302
303
304
305
306
307
308
309
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
310
	s.seqsSem.Release(1)
311
312
313
314
315
}

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

316
	// Logically these batches are used only within the context of processBatch
317
	// but it is better for performance to allocate them once here
Jesse Gross's avatar
Jesse Gross committed
318
319
320
321
	tokenBatch, err := llama.NewBatch(s.batchSize, len(s.seqs), 0)
	if err != nil {
		panic(err)
	}
322
323
	defer tokenBatch.Free()

324
325
326
	var embedBatch *llama.Batch
	embedBatchSize := s.image.BatchSize(s.batchSize)
	if embedBatchSize != 0 {
Jesse Gross's avatar
Jesse Gross committed
327
328
329
330
		embedBatch, err = llama.NewBatch(embedBatchSize, len(s.seqs), s.image.EmbedSize(s.lc))
		if err != nil {
			panic(err)
		}
331
332
333
334
		defer embedBatch.Free()
	} else {
		embedBatch = &llama.Batch{}
	}
335
336
337
338
339
340

	for {
		select {
		case <-ctx.Done():
			return
		default:
341
342
343
344
345
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

346
347
348
349
350
351
352
353
354
355
356
357
358
			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
359
func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch) error {
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
	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
378
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
379
			s.removeSequence(seqIdx, llm.DoneReasonLength)
380
381
382
383
			continue
		}

		for i, input := range seq.inputs {
384
385
386
387
			if len(seq.cache.Inputs)+len(seq.pendingInputs)+1 > s.cache.numCtx {
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
388
389
390
391
392
393
394
395
396
						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
						}
397
					}
398
399
400
401
402
				} else {
					break
				}
			}

403
404
405
406
407
408
409
410
411
412
413
414
			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
				}
415
			} else if embedding != batch.IsEmbedding() {
416
417
418
419
				s.nextSeq = seqIdx
				break
			}

420
			if i >= batch.Size() {
421
422
423
				break
			}

424
425
			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)
426
427
			seq.iBatch = batch.NumTokens() - 1
		}
428
429

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

	if batch == nil || batch.NumTokens() == 0 {
433
		return nil
434
435
436
437
	}

	err := s.lc.Decode(batch)
	if err != nil {
438
		return fmt.Errorf("failed to decode batch: %w", err)
439
440
441
442
443
444
445
	}

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

446
447
448
449
450
451
		// 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{}
		}

452
453
454
455
456
457
458
459
460
461
462
463
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
			continue
		}

		seq.numDecoded += 1
		if seq.numDecoded == 1 {
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
464
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
465
466
467
468
469
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
470
			s.removeSequence(i, llm.DoneReasonStop)
471
472
473
474
			continue
		}

		// sample a token
475
476
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
477
478
479
480
481
482
483
484
485
486
		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

487
			s.removeSequence(i, llm.DoneReasonStop)
488
489
490
491
492
493
494
495
			continue
		}

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

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

Jesse Gross's avatar
Jesse Gross committed
496
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
497
498
499
500
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

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

518
			s.removeSequence(i, llm.DoneReasonStop)
519
520
521
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
522
		if common.ContainsStopSuffix(sequence, seq.stop) {
523
524
525
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
526
		if common.IncompleteUnicode(sequence) {
527
528
529
530
			continue
		}

		if !flushPending(seq) {
531
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
532
533
		}
	}
534
535

	return nil
536
537
538
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
539
	var req llm.CompletionRequest
540
541
542
543
544
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

545
546
547
548
549
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

550
551
552
553
554
555
556
557
558
559
	// 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
	}

560
561
562
563
564
565
566
567
568
569
570
571
572
573
	// 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,
	}
574
575

	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
576
577
578
		numPredict:     req.Options.NumPredict,
		stop:           req.Options.Stop,
		numKeep:        req.Options.NumKeep,
579
580
581
582
583
584
585
586
		samplingParams: &samplingParams,
		embedding:      false,
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

587
	// Ensure there is a place to put the sequence, released when removed from s.seqs
588
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
589
590
591
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
592
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
593
		}
594
595
596
		return
	}

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

609
610
			s.seqs[i] = seq
			s.cond.Signal()
611
			found = true
612
613
614
615
616
			break
		}
	}
	s.mu.Unlock()

617
	if !found {
618
		s.seqsSem.Release(1)
619
620
621
622
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

623
624
625
626
627
628
629
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
630
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
631
632
633
634
635
636
637
638
639
					Content: content,
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
640
641
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
642
					DoneReason:         seq.doneReason,
643
644
645
646
					PromptEvalCount:    seq.numPromptInputs,
					PromptEvalDuration: seq.startGenerationTime.Sub(seq.startProcessingTime),
					EvalCount:          seq.numDecoded,
					EvalDuration:       time.Since(seq.startGenerationTime),
647
648
649
650
651
652
653
654
655
656
657
				}); 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) {
658
	var req llm.EmbeddingRequest
659
660
661
662
663
664
665
666
667
668
669
670
671
	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
	}

672
	// Ensure there is a place to put the sequence, released when removed from s.seqs
673
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
674
675
676
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
677
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
678
		}
679
680
681
		return
	}

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

701
	if !found {
702
		s.seqsSem.Release(1)
703
704
705
706
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

707
708
	embedding := <-seq.embedding

709
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
710
711
712
713
714
715
716
717
		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")
718
719
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
720
721
722
723
724
725
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

726
727
728
729
730
731
732
733
734
735
736
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, ", ")
}

737
738
739
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
740
	lpath multiLPath,
741
742
	ppath string,
	kvSize int,
743
	kvCacheType string,
744
745
746
747
	flashAttention bool,
	threads int,
	multiUserCache bool,
) {
748
749
750
751
752
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
753

754
	ctxParams := llama.NewContextParams(kvSize, s.batchSize*s.parallel, s.parallel, threads, flashAttention, kvCacheType)
755
756
757
758
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
759

760
761
762
763
764
765
	if lpath.String() != "" {
		for _, path := range lpath {
			err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
			if err != nil {
				panic(err)
			}
766
767
768
769
		}
	}

	if ppath != "" {
770
		var err error
771
		s.image, err = NewImageContext(s.lc, ppath)
772
773
774
		if err != nil {
			panic(err)
		}
775
776
	}

777
778
779
780
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
781

782
	s.status = llm.ServerStatusReady
783
784
785
	s.ready.Done()
}

786
787
788
789
790
791
792
793
794
795
796
797
798
func Execute(args []string) error {
	fs := flag.NewFlagSet("runner", flag.ExitOnError)
	mpath := fs.String("model", "", "Path to model binary file")
	ppath := fs.String("mmproj", "", "Path to projector binary file")
	parallel := fs.Int("parallel", 1, "Number of sequences to handle simultaneously")
	batchSize := fs.Int("batch-size", 512, "Batch size")
	nGpuLayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGpu := fs.Int("main-gpu", 0, "Main GPU")
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
	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")
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
799
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
800
801
802
	noMmap := fs.Bool("no-mmap", false, "do not memory-map model (slower load but may reduce pageouts if not using mlock)")
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
803

804
	var lpaths multiLPath
805
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
806

807
808
809
810
811
812
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
813
	}
814
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
815
	slog.Info("starting go runner")
816
817

	llama.BackendInit()
818
819
820
821
822

	server := &Server{
		batchSize: *batchSize,
		parallel:  *parallel,
		seqs:      make([]*Sequence, *parallel),
823
		seqsSem:   semaphore.NewWeighted(int64(*parallel)),
824
		status:    llm.ServerStatusLoadingModel,
825
826
827
828
	}

	var tensorSplitFloats []float32
	if *tensorSplit != "" {
829
830
831
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
832
			f, _ := strconv.ParseFloat(s, 32)
833
			tensorSplitFloats[i] = float32(f)
834
835
836
837
838
839
		}
	}

	params := llama.ModelParams{
		NumGpuLayers: *nGpuLayers,
		MainGpu:      *mainGpu,
840
		UseMmap:      !*noMmap && lpaths.String() == "",
841
842
843
844
845
846
847
		TensorSplit:  tensorSplitFloats,
		Progress: func(progress float32) {
			server.progress = progress
		},
	}

	server.ready.Add(1)
848
	go server.loadModel(params, *mpath, lpaths, *ppath, *kvSize, *kvCacheType, *flashAttention, *threads, *multiUserCache)
849
850
851
852

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

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

855
856
857
858
859
860
	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)
861
		return err
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
	}
	defer listener.Close()

	mux := http.NewServeMux()
	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)
877
		return err
878
879
	}

880
	return nil
881
}