runner.go 24 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package ollamarunner
2
3
4
5
6
7
8

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
9
	"hash/maphash"
10
11
12
13
14
15
16
17
18
19
20
21
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
22
	"unicode/utf8"
23

24
25
	"golang.org/x/sync/semaphore"

26
	"github.com/ollama/ollama/api"
27
	"github.com/ollama/ollama/ml"
Jesse Gross's avatar
Jesse Gross committed
28
29
30
31
32
	"github.com/ollama/ollama/model"
	"github.com/ollama/ollama/runner/common"
	"github.com/ollama/ollama/sample"

	_ "github.com/ollama/ollama/model/models"
33
34
35
)

type Sequence struct {
36
37
38
39
	// ctx for allocating tensors that last the lifetime of the sequence, such as
	// multimodal embeddings
	ctx ml.Context

40
41
42
43
	// batch index
	iBatch int

	// prompt inputs left to evaluate
44
	inputs []model.Input
45

Jesse Gross's avatar
Jesse Gross committed
46
	// inputs that have been added to a batch but not yet submitted to Forward
47
	pendingInputs []model.Input
48

49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
	// 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

64
65
	// sampler with transforms to run on generated logits
	sampler sample.Sampler
66
67
68
69
70
71
72
73

	// 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
Jesse Gross's avatar
Jesse Gross committed
74
	numKeep int32
75
76
77
78
79
80
81
82
83

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

	doneReason string

	// Metrics
	startProcessingTime time.Time
	startGenerationTime time.Time
Jesse Gross's avatar
Jesse Gross committed
84
	numPredicted        int
85
86
87
88
	numPromptInputs     int
}

type NewSequenceParams struct {
Jesse Gross's avatar
Jesse Gross committed
89
90
91
	numPredict int
	stop       []string
	numKeep    int32
92
	sampler    sample.Sampler
Jesse Gross's avatar
Jesse Gross committed
93
	embedding  bool
94
95
96
97
98
99
}

func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequenceParams) (*Sequence, error) {
	s.ready.Wait()

	startTime := time.Now()
100
	ctx := s.model.Backend().NewContext()
101

102
	inputs, err := s.inputs(ctx, prompt, images)
103
104
105
106
107
108
109
	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 {
Jesse Gross's avatar
Jesse Gross committed
110
		params.numKeep = int32(len(inputs))
111
112
	}

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

Jesse Gross's avatar
Jesse Gross committed
116
117
	if int32(len(inputs)) > s.cache.numCtx {
		discard := int32(len(inputs)) - s.cache.numCtx
118
		newInputs := inputs[:params.numKeep]
119
120
121
		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))
122
		inputs = newInputs
123
124
	}

Jesse Gross's avatar
Jesse Gross committed
125
	// TODO(jessegross): Ingest cached history for grammar
126
127

	return &Sequence{
128
		ctx:                 ctx,
129
130
131
132
133
134
135
136
		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),
137
		sampler:             params.sampler,
138
139
140
141
142
143
144
145
		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
Jesse Gross's avatar
Jesse Gross committed
146
// decoding images
147
148
func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]model.Input, error) {
	var inputs []model.Input
149
150
151
	var parts []string
	var matches [][]string

152
	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
153

154
155
156
157
158
159
160
161
162
	if visionModel {
		re := regexp.MustCompile(`\[img-(\d+)\]`)
		parts = re.Split(prompt, -1)
		matches = re.FindAllStringSubmatch(prompt, -1)
	} else {
		parts = []string{prompt}
	}

	postTokenize := false
163
164
	for i, part := range parts {
		// text - tokenize
165
		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
166
167
168
		if err != nil {
			return nil, err
		}
169

170
		for _, t := range tokens {
171
			inputs = append(inputs, model.Input{Token: t})
172
173
		}

Jesse Gross's avatar
Jesse Gross committed
174
		// image - decode and store
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
		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)
			}

190
			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
191
192
193
194
			if err != nil {
				return nil, err
			}

195
196
197
198
199
200
201
202
203
204
205
206
207
208
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

			inputs = append(inputs, model.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
			postTokenize = true
		}
	}

	if visionModel && postTokenize {
		var err error
		inputs, err = multimodalProcessor.PostTokenize(ctx, inputs)
		if err != nil {
			return nil, err
209
210
211
212
213
214
215
		}
	}

	return inputs, nil
}

type Server struct {
216
217
218
219
220
	// is the server ready to process requests?
	// protects access to model and image
	ready sync.WaitGroup

	// loaded model
Jesse Gross's avatar
Jesse Gross committed
221
	model model.Model
222

223
224
225
226
227
228
229
230
231
232
	// status for external health reporting - loading, ready to serve, etc.
	status ServerStatus

	// 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)
233
	// TODO (jmorganca): make this n_batch
234
235
	batchSize int

236
237
238
239
240
241
242
243
	// 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

	// the list of simultaneous sequences being evaluated
244
245
	seqs []*Sequence

246
247
248
249
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

250
251
252
253
254
	// KV cache
	cache *InputCache

	// next sequence for prompt processing to avoid starvation
	nextSeq int
255
256
257
258

	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
259
260
261
262
263
264
265
266
267
268
269
270
}

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

func flushPending(seq *Sequence) bool {
271
272
273
274
275
276
277
278
279
280
281
	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]
282
283
	}

284
285
286
287
288
289
290
291
292
293
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
294
295
296
297
298
299
300
301
302
303
}

func (s *Server) removeSequence(seqIndex int, reason string) {
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
304
	seq.ctx.Close()
305
	s.seqs[seqIndex] = nil
306
	s.seqsSem.Release(1)
307
308
309
310
311
312
313
314
315
316
}

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

	for {
		select {
		case <-ctx.Done():
			return
		default:
Jesse Gross's avatar
Jesse Gross committed
317
			err := s.processBatch()
318
319
320
			if err != nil {
				panic(err)
			}
321
322
323
324
		}
	}
}

Jesse Gross's avatar
Jesse Gross committed
325
func (s *Server) processBatch() error {
326
327
328
329
330
331
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

Jesse Gross's avatar
Jesse Gross committed
332
	var options model.Options
333
334
335
336
337
338
339
340
341
342
343

	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
344
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
345
346
347
348
			s.removeSequence(seqIdx, "limit")
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
349
350
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
351
			seq.cache.Inputs = []model.Input{}
Jesse Gross's avatar
Jesse Gross committed
352
353
		}

354
		for i, input := range seq.inputs {
Jesse Gross's avatar
Jesse Gross committed
355
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
356
357
358
359
360
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
						return err
					}
361
362
363
364
365
				} else {
					break
				}
			}

Jesse Gross's avatar
Jesse Gross committed
366
			if i >= s.batchSize {
367
368
369
				break
			}

370
371
372
373
374
375
376
377
378
			// TODO(jessegross): This is a workaround for generating an attention mask and also providing a hint
			// to the encoder cache.
			//
			// Break the batch when switching from text to images so that images are always at the beginning.
			if input.Multimodal != nil && !(len(seq.pendingInputs) == 0 ||
				(len(options.Multimodal) > 0 && options.Multimodal[len(options.Multimodal)-1].Index == len(options.Inputs)-1)) {
				s.nextSeq = seqIdx
				break
			}
Jesse Gross's avatar
Jesse Gross committed
379

380
381
382
			options.Inputs = append(options.Inputs, input.Token)
			if input.Multimodal != nil {
				options.Multimodal = append(options.Multimodal, model.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: input.Multimodal})
383
384
			}

Jesse Gross's avatar
Jesse Gross committed
385
386
387
388
389
390
391
			options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
			options.Sequences = append(options.Sequences, seq.cache.Id)

			seq.iBatch = len(options.Outputs)
			if i+1 == len(seq.inputs) {
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
392
			seq.pendingInputs = append(seq.pendingInputs, input)
393
		}
394
395

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
396
397
	}

Jesse Gross's avatar
Jesse Gross committed
398
	if len(options.Inputs) == 0 {
399
		return nil
400
401
	}

Jesse Gross's avatar
Jesse Gross committed
402
403
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
404

Jesse Gross's avatar
Jesse Gross committed
405
	modelOutput, err := model.Forward(ctx, s.model, options)
406
	if err != nil {
407
		return fmt.Errorf("failed to decode batch: %w", err)
408
409
	}

410
	logits := modelOutput.Floats()
411

412
413
414
415
416
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
417
		// After calling Forward, pending inputs are now in the cache
418
419
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
420
			seq.pendingInputs = []model.Input{}
421
422
		}

423
424
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
425
426
427
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
428
429
430
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
431
432
		seq.numPredicted++
		if seq.numPredicted == 1 {
433
434
435
436
437
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
438
			// TODO(jessegross): Embedding support
439
440
441
442
			// s.removeSequence(i, "")
			// continue

			panic("generation of embedding outputs not yet supported")
443
444
445
		}

		// sample a token
446
447
448
		vocabSize := len(logits) / len(options.Outputs)

		token, err := seq.sampler.Sample(logits[seq.iBatch*vocabSize : (seq.iBatch+1)*vocabSize])
Jesse Gross's avatar
Jesse Gross committed
449
		if err != nil {
450
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
451
		}
452
453

		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
454
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
455
456
457
458
459
460
461
462
			// TODO (jmorganca): we should send this back
			// as it's important for the /api/generate context
			// seq.responses <- piece

			s.removeSequence(i, "stop")
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
463
464
465
466
467
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
			return err
		}

468
		seq.inputs = []model.Input{{Token: token}}
469
470
471
472

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

Jesse Gross's avatar
Jesse Gross committed
473
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
474
475
476
477
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
478
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
			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]
494
495
496
497
498

			s.removeSequence(i, "stop")
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
499
		if common.ContainsStopSuffix(sequence, seq.stop) {
500
501
502
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
503
		if common.IncompleteUnicode(sequence) {
504
505
506
507
508
509
510
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
511
512

	return nil
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
}

// TODO (jmorganca): use structs from the api package to avoid duplication
// this way the api acts as a proxy instead of using a different api for the
// runner
type Options struct {
	api.Runner

	NumKeep          int      `json:"n_keep"`
	Seed             int      `json:"seed"`
	NumPredict       int      `json:"n_predict"`
	TopK             int      `json:"top_k"`
	TopP             float32  `json:"top_p"`
	MinP             float32  `json:"min_p"`
	TypicalP         float32  `json:"typical_p"`
	RepeatLastN      int      `json:"repeat_last_n"`
	Temperature      float32  `json:"temperature"`
	RepeatPenalty    float32  `json:"repeat_penalty"`
	PresencePenalty  float32  `json:"presence_penalty"`
	FrequencyPenalty float32  `json:"frequency_penalty"`
	Mirostat         int      `json:"mirostat"`
	MirostatTau      float32  `json:"mirostat_tau"`
	MirostatEta      float32  `json:"mirostat_eta"`
	Stop             []string `json:"stop"`
}

type ImageData struct {
540
541
542
	Data          []byte `json:"data"`
	ID            int    `json:"id"`
	AspectRatioID int    `json:"aspect_ratio_id"`
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
}

type CompletionRequest struct {
	Prompt      string      `json:"prompt"`
	Images      []ImageData `json:"image_data"`
	Grammar     string      `json:"grammar"`
	CachePrompt bool        `json:"cache_prompt"`

	Options
}

type Timings struct {
	PredictedN  int     `json:"predicted_n"`
	PredictedMS float64 `json:"predicted_ms"`
	PromptN     int     `json:"prompt_n"`
	PromptMS    float64 `json:"prompt_ms"`
}

type CompletionResponse struct {
	Content string `json:"content"`
	Stop    bool   `json:"stop"`

	Model        string  `json:"model,omitempty"`
	Prompt       string  `json:"prompt,omitempty"`
	StoppedLimit bool    `json:"stopped_limit,omitempty"`
	PredictedN   int     `json:"predicted_n,omitempty"`
	PredictedMS  float64 `json:"predicted_ms,omitempty"`
	PromptN      int     `json:"prompt_n,omitempty"`
	PromptMS     float64 `json:"prompt_ms,omitempty"`

	Timings Timings `json:"timings"`
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
	var req CompletionRequest
	req.Options = Options(api.DefaultOptions())
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

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

594
595
596
597
598
599
600
601
	sampler := sample.NewSampler(
		req.Temperature,
		req.TopK,
		req.TopP,
		req.MinP,
		req.Seed,
	)

602
603
604
605
	if req.Grammar != "" {
		panic("grammars are not yet supported")
	}

606
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
Jesse Gross's avatar
Jesse Gross committed
607
608
609
		numPredict: req.NumPredict,
		stop:       req.Stop,
		numKeep:    int32(req.NumKeep),
610
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
611
		embedding:  false,
612
613
614
615
616
617
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

618
	// Ensure there is a place to put the sequence, released when removed from s.seqs
619
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
620
621
622
623
624
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
			slog.Error("Failed to acquire semaphore", "error", err)
		}
625
626
627
		return
	}

628
	s.mu.Lock()
629
	found := false
630
631
	for i, sq := range s.seqs {
		if sq == nil {
632
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
633
634
635
636
637
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
638

639
640
			s.seqs[i] = seq
			s.cond.Signal()
641
			found = true
642
643
644
645
646
			break
		}
	}
	s.mu.Unlock()

647
648
649
650
651
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

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

				flusher.Flush()
			} else {
				// Send the final response
				if err := json.NewEncoder(w).Encode(&CompletionResponse{
					Stop:         true,
					StoppedLimit: seq.doneReason == "limit",
					Timings: Timings{
						PromptN:     seq.numPromptInputs,
						PromptMS:    float64(seq.startGenerationTime.Sub(seq.startProcessingTime).Milliseconds()),
Jesse Gross's avatar
Jesse Gross committed
676
						PredictedN:  seq.numPredicted,
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
						PredictedMS: float64(time.Since(seq.startGenerationTime).Milliseconds()),
					},
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
				}

				return
			}
		}
	}
}

type EmbeddingRequest struct {
	Content     string `json:"content"`
	CachePrompt bool   `json:"cache_prompt"`
}

type EmbeddingResponse struct {
	Embedding []float32 `json:"embedding"`
}

func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
	var req EmbeddingRequest
	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")

	slog.Debug("embedding request", "content", req.Content)

	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
	}

715
	// Ensure there is a place to put the sequence, released when removed from s.seqs
716
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
717
718
719
720
721
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embeddings request due to client closing the connection")
		} else {
			slog.Error("Failed to acquire semaphore", "error", err)
		}
722
723
724
		return
	}

725
	s.mu.Lock()
726
	found := false
727
728
	for i, sq := range s.seqs {
		if sq == nil {
729
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
730
731
732
733
734
735
736
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
			s.seqs[i] = seq
			s.cond.Signal()
737
			found = true
738
739
740
741
742
			break
		}
	}
	s.mu.Unlock()

743
744
745
746
747
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
	embedding := <-seq.embedding

	if err := json.NewEncoder(w).Encode(&EmbeddingResponse{
		Embedding: embedding,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

type HealthResponse struct {
	Status   string  `json:"status"`
	Progress float32 `json:"progress"`
}

type ServerStatus int

const (
	ServerStatusReady ServerStatus = iota
	ServerStatusLoadingModel
	ServerStatusError
)

func (s ServerStatus) ToString() string {
	switch s {
	case ServerStatusReady:
		return "ok"
	case ServerStatusLoadingModel:
		return "loading model"
	default:
		return "server error"
	}
}

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

791
792
793
794
795
796
797
798
799
800
801
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, ", ")
}

802
803
func (s *Server) loadModel(
	mpath string,
804
	params ml.BackendParams,
805
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
806
	parallel int,
807
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
808
	kvSize int,
809
810
	multiUserCache bool,
) {
811
	var err error
812
	s.model, err = model.New(mpath, params)
813
814
815
	if err != nil {
		panic(err)
	}
816

Jesse Gross's avatar
Jesse Gross committed
817
	// TODO(jessegross): LoRA loading
818
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
819
		panic("loras are not yet implemented")
820
821
	}

Jesse Gross's avatar
Jesse Gross committed
822
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
823
824
825
	if err != nil {
		panic(err)
	}
826

Jesse Gross's avatar
Jesse Gross committed
827
828
829
830
831
832
833
834
835
	if !s.cache.enabled && parallel > 1 {
		parallel = 1
		slog.Warn("model does not support caching, disabling parallel processing")
	}

	s.parallel = parallel
	s.seqs = make([]*Sequence, s.parallel)
	s.seqsSem = semaphore.NewWeighted(int64(s.parallel))

836
837
838
839
	s.status = ServerStatusReady
	s.ready.Done()
}

840
841
842
843
844
func Execute(args []string) error {
	fs := flag.NewFlagSet("runner", flag.ExitOnError)
	mpath := fs.String("model", "", "Path to model binary file")
	parallel := fs.Int("parallel", 1, "Number of sequences to handle simultaneously")
	batchSize := fs.Int("batch-size", 512, "Batch size")
845
846
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
847
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
848
849
850
	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")
851
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
852
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
853
854
	_ = fs.Bool("no-mmap", false, "do not memory-map model (slower load but may reduce pageouts if not using mlock)")
	_ = fs.Bool("mlock", false, "force system to keep model in RAM rather than swapping or compressing")
855
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
856
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
857

858
	var lpaths multiLPath
859
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
860

861
862
863
864
865
866
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
	}
	level := slog.LevelInfo
	if *verbose {
		level = slog.LevelDebug
	}
	handler := slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level:     level,
		AddSource: true,
		ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
			if attr.Key == slog.SourceKey {
				source := attr.Value.Any().(*slog.Source)
				source.File = filepath.Base(source.File)
			}
			return attr
		},
	})
	slog.SetDefault(slog.New(handler))
Jesse Gross's avatar
Jesse Gross committed
884
	slog.Info("starting ollama engine")
885
886
887
888
889
890

	server := &Server{
		batchSize: *batchSize,
		status:    ServerStatusLoadingModel,
	}

Jesse Gross's avatar
Jesse Gross committed
891
892
893
894
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

895
	var tensorSplitFloats []float32
896
	if *tensorSplit != "" {
897
898
899
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
900
			f, _ := strconv.ParseFloat(s, 32)
901
			tensorSplitFloats[i] = float32(f)
902
		}
903
904
905
	}

	params := ml.BackendParams{
906
907
908
909
910
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
911
	}
912
913

	server.ready.Add(1)
914
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
915
916
917
918

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

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

921
922
923
924
925
926
	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)
927
		return err
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
	}
	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)
943
		return err
944
945
	}

946
	return nil
947
}