"docs/get_started/install.md" did not exist on "fadf18fdd58f6c894f3204d937f5be40bbb67bc2"
runner.go 22.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
	doneReason llm.DoneReason
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

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

99
func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
	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)
	}

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

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

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

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

	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
163
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
164
	var inputs []input
165
166
167
168
169
170
171
172
173
174
	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}
	}
175
176
177

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

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

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

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

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

	return inputs, nil
}

type Server struct {
Jesse Gross's avatar
Jesse Gross committed
218
219
220
221
222
223
	// 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

224
225
226
227
228
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
229
230
	model *llama.Model

231
	// image model context for multi-modal models
232
	image *ImageContext
233

234
	// status for external health reporting - loading, ready to serve, etc.
235
	status llm.ServerStatus
236
237
238
239
240
241
242
243

	// 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)
244
	// TODO (jmorganca): make this n_batch
245
246
	batchSize int

247
248
249
250
251
252
253
254
255
	// 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
256

257
	// the list of simultaneous sequences being evaluated
258
259
	seqs []*Sequence

260
261
262
263
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
	// 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 {
281
282
283
284
285
286
287
288
289
290
291
	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]
292
293
	}

294
295
296
297
298
299
300
301
302
303
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
304
305
}

306
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
307
308
309
310
311
312
313
314
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
315
	s.seqsSem.Release(1)
316
317
318
319
320
}

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

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

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

	for {
		select {
		case <-ctx.Done():
			return
		default:
346
347
348
349
350
			err := s.processBatch(tokenBatch, embedBatch)
			if err != nil {
				panic(err)
			}

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

		for i, input := range seq.inputs {
389
390
391
392
			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 {
393
394
395
396
397
398
399
400
401
						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
						}
402
					}
403
404
405
406
407
				} else {
					break
				}
			}

408
409
410
411
412
413
414
415
416
417
418
419
			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
				}
420
			} else if embedding != batch.IsEmbedding() {
421
422
423
424
				s.nextSeq = seqIdx
				break
			}

425
			if i >= batch.Size() {
426
427
428
				break
			}

429
430
			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)
431
432
			seq.iBatch = batch.NumTokens() - 1
		}
433
434

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
435
436
437
	}

	if batch == nil || batch.NumTokens() == 0 {
438
		return nil
439
440
441
442
	}

	err := s.lc.Decode(batch)
	if err != nil {
443
		return fmt.Errorf("failed to decode batch: %w", err)
444
445
446
447
448
449
450
	}

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

451
452
453
454
455
456
		// 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{}
		}

457
458
459
460
461
462
463
464
465
466
467
468
		// 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 {
469
			embed := s.lc.GetEmbeddingsSeq(seq.cache.Id)
470
471
472
473
474
			if embed == nil {
				embed = s.lc.GetEmbeddingsIth(seq.iBatch)
			}

			seq.embedding <- embed
475
			s.removeSequence(i, llm.DoneReasonStop)
476
477
478
479
			continue
		}

		// sample a token
480
481
		token := seq.samplingCtx.Sample(s.lc, seq.iBatch)
		seq.samplingCtx.Accept(token, true)
482
483
484
485
486
487
488
489
490
491
		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

492
			s.removeSequence(i, llm.DoneReasonStop)
493
494
495
496
497
498
499
500
			continue
		}

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

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

Jesse Gross's avatar
Jesse Gross committed
501
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
502
503
504
505
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

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

523
			s.removeSequence(i, llm.DoneReasonStop)
524
525
526
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
527
		if common.ContainsStopSuffix(sequence, seq.stop) {
528
529
530
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
531
		if common.IncompleteUnicode(sequence) {
532
533
534
535
			continue
		}

		if !flushPending(seq) {
536
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
537
538
		}
	}
539
540

	return nil
541
542
543
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
544
	var req llm.CompletionRequest
545
546
547
548
549
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

550
551
552
553
554
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

555
556
557
558
559
560
561
562
563
564
	// 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
	}

565
566
567
568
569
570
571
572
573
574
575
576
577
578
	// 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,
	}
579
580

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

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

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

614
615
			s.seqs[i] = seq
			s.cond.Signal()
616
			found = true
617
618
619
620
621
			break
		}
	}
	s.mu.Unlock()

622
	if !found {
623
		s.seqsSem.Release(1)
624
625
626
627
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

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

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

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

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

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

712
713
	embedding := <-seq.embedding

714
	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
715
716
717
718
719
720
721
722
		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")
723
724
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
725
726
727
728
729
730
		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
731
732
// 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.
733
734
735
func (s *Server) loadModel(
	params llama.ModelParams,
	mpath string,
Jesse Gross's avatar
Jesse Gross committed
736
	lpath []string,
737
738
	ppath string,
	kvSize int,
739
	kvCacheType string,
740
741
742
743
	flashAttention bool,
	threads int,
	multiUserCache bool,
) {
744
745
746
747
748
	var err error
	s.model, err = llama.LoadModelFromFile(mpath, params)
	if err != nil {
		panic(err)
	}
749

750
	ctxParams := llama.NewContextParams(kvSize, s.batchSize*s.parallel, s.parallel, threads, flashAttention, kvCacheType)
751
752
753
754
	s.lc, err = llama.NewContextWithModel(s.model, ctxParams)
	if err != nil {
		panic(err)
	}
755

Jesse Gross's avatar
Jesse Gross committed
756
757
758
759
	for _, path := range lpath {
		err := s.model.ApplyLoraFromFile(s.lc, path, 1.0, threads)
		if err != nil {
			panic(err)
760
761
762
763
		}
	}

	if ppath != "" {
764
		var err error
765
		s.image, err = NewImageContext(s.lc, ppath)
766
767
768
		if err != nil {
			panic(err)
		}
769
770
	}

771
772
773
774
	s.cache, err = NewInputCache(s.lc, kvSize, s.parallel, multiUserCache)
	if err != nil {
		panic(err)
	}
775

776
	s.status = llm.ServerStatusReady
777
778
779
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
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
837
838
839
840
841
842
843
844
845
846
847
848
849
// 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 {
				if gpuIDs[i] == layers.ID {
					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
	}
}

850
851
852
853
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")
854
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
855

856
857
858
859
860
861
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
862
	}
863
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
864
	slog.Info("starting go runner")
865
866

	llama.BackendInit()
867
868

	server := &Server{
Jesse Gross's avatar
Jesse Gross committed
869
870
		modelPath: *mpath,
		status:    llm.ServerStatusLaunched,
871
872
873
874
875
876
877
	}

	server.ready.Add(1)

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

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

880
881
882
883
884
885
	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)
886
		return err
887
888
889
890
	}
	defer listener.Close()

	mux := http.NewServeMux()
Jesse Gross's avatar
Jesse Gross committed
891
	mux.HandleFunc("POST /load", server.load)
892
893
894
895
896
897
898
899
900
901
902
	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)
903
		return err
904
905
	}

906
	return nil
907
}