runner.go 23.5 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
	"github.com/ollama/ollama/model"
29
	"github.com/ollama/ollama/model/input"
Jesse Gross's avatar
Jesse Gross committed
30
31
32
33
	"github.com/ollama/ollama/runner/common"
	"github.com/ollama/ollama/sample"

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

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

41
42
43
44
	// batch index
	iBatch int

	// prompt inputs left to evaluate
45
	inputs []input.Input
46

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

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

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

	// 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
75
	numKeep int32
76
77
78
79
80
81
82
83
84

	// 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
85
	numPredicted        int
86
87
88
89
	numPromptInputs     int
}

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

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

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

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

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

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

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

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

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

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

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

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

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

196
197
198
199
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

200
			inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
201
202
203
204
205
206
207
208
209
			postTokenize = true
		}
	}

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

	return inputs, nil
}

type Server struct {
217
218
219
220
221
	// 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
222
	model model.Model
223

224
225
226
227
228
229
230
231
232
233
	// 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)
234
	// TODO (jmorganca): make this n_batch
235
236
	batchSize int

237
238
239
240
241
242
243
244
	// 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
245
246
	seqs []*Sequence

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

251
252
253
	// KV cache
	cache *InputCache

254
255
256
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
257
258
259
260
261
262
263
264
265
266
267
268
}

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

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

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

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

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
302
	seq.ctx.Close()
303
	s.seqs[seqIndex] = nil
304
	s.seqsSem.Release(1)
305
306
307
308
309
310
311
312
313
314
}

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

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

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

330
	var options input.Options
331

332
	for i, seq := range s.seqs {
333
334
335
336
337
		if seq == nil {
			continue
		}

		// if past the num predict limit
338
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
339
			s.removeSequence(i, "limit")
340
341
342
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
343
344
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
345
			seq.cache.Inputs = []input.Input{}
Jesse Gross's avatar
Jesse Gross committed
346
347
		}

348
		for j, inp := range seq.inputs {
Jesse Gross's avatar
Jesse Gross committed
349
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
350
351
352
353
354
				if len(seq.pendingInputs) == 0 {
					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
					if err != nil {
						return err
					}
355
356
357
358
359
				} else {
					break
				}
			}

360
			if j >= s.batchSize {
361
362
				break
			}
Jesse Gross's avatar
Jesse Gross committed
363

364
365
366
			options.Inputs = append(options.Inputs, inp.Token)
			if inp.Multimodal != nil {
				options.Multimodal = append(options.Multimodal, input.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: inp.Multimodal})
367
368
			}

Jesse Gross's avatar
Jesse Gross committed
369
370
371
372
			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)
373
			if j+1 == len(seq.inputs) {
Jesse Gross's avatar
Jesse Gross committed
374
375
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
376
			seq.pendingInputs = append(seq.pendingInputs, inp)
377
		}
378
379

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
380
381
	}

Jesse Gross's avatar
Jesse Gross committed
382
	if len(options.Inputs) == 0 {
383
		return nil
384
385
	}

Jesse Gross's avatar
Jesse Gross committed
386
387
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
388

Jesse Gross's avatar
Jesse Gross committed
389
	modelOutput, err := model.Forward(ctx, s.model, options)
390
	if err != nil {
391
		return fmt.Errorf("failed to decode batch: %w", err)
392
393
	}

394
	logits := modelOutput.Floats()
395

396
397
398
399
400
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
401
		// After calling Forward, pending inputs are now in the cache
402
403
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
404
			seq.pendingInputs = []input.Input{}
405
406
		}

407
408
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
409
410
411
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
412
413
414
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
415
416
		seq.numPredicted++
		if seq.numPredicted == 1 {
417
418
419
420
421
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
422
			// TODO(jessegross): Embedding support
423
424
425
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
426
427
428
		}

		// sample a token
429
430
431
		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
432
		if err != nil {
433
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
434
		}
435
436

		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
437
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
438
439
440
441
442
443
444
445
			// 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
446
447
448
449
450
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
			return err
		}

451
		seq.inputs = []input.Input{{Token: token}}
452
453
454
455

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

Jesse Gross's avatar
Jesse Gross committed
456
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
457
458
459
460
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
461
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
			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]
477
478
479
480
481

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

Jesse Gross's avatar
Jesse Gross committed
482
		if common.ContainsStopSuffix(sequence, seq.stop) {
483
484
485
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
486
		if common.IncompleteUnicode(sequence) {
487
488
489
490
491
492
493
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
494
495

	return nil
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
}

// 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 {
523
524
525
	Data          []byte `json:"data"`
	ID            int    `json:"id"`
	AspectRatioID int    `json:"aspect_ratio_id"`
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
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
}

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
	}

577
578
579
580
581
582
583
584
	sampler := sample.NewSampler(
		req.Temperature,
		req.TopK,
		req.TopP,
		req.MinP,
		req.Seed,
	)

585
586
587
588
	if req.Grammar != "" {
		panic("grammars are not yet supported")
	}

589
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
Jesse Gross's avatar
Jesse Gross committed
590
591
592
		numPredict: req.NumPredict,
		stop:       req.Stop,
		numKeep:    int32(req.NumKeep),
593
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
594
		embedding:  false,
595
596
597
598
599
600
	})
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

601
	// Ensure there is a place to put the sequence, released when removed from s.seqs
602
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
603
604
605
606
607
		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)
		}
608
609
610
		return
	}

611
	s.mu.Lock()
612
	found := false
613
614
	for i, sq := range s.seqs {
		if sq == nil {
615
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
616
617
618
619
620
			if err != nil {
				s.mu.Unlock()
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
621

622
623
			s.seqs[i] = seq
			s.cond.Signal()
624
			found = true
625
626
627
628
629
			break
		}
	}
	s.mu.Unlock()

630
631
632
633
634
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
	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
659
						PredictedN:  seq.numPredicted,
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
						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
	}

698
	// Ensure there is a place to put the sequence, released when removed from s.seqs
699
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
700
701
702
703
704
		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)
		}
705
706
707
		return
	}

708
	s.mu.Lock()
709
	found := false
710
711
	for i, sq := range s.seqs {
		if sq == nil {
712
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
713
714
715
716
717
718
719
			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()
720
			found = true
721
722
723
724
725
			break
		}
	}
	s.mu.Unlock()

726
727
728
729
730
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
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
	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)
	}
}

774
775
776
777
778
779
780
781
782
783
784
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, ", ")
}

785
786
func (s *Server) loadModel(
	mpath string,
787
	params ml.BackendParams,
788
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
789
	parallel int,
790
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
791
	kvSize int,
792
793
	multiUserCache bool,
) {
794
	var err error
795
	s.model, err = model.New(mpath, params)
796
797
798
	if err != nil {
		panic(err)
	}
799

Jesse Gross's avatar
Jesse Gross committed
800
	// TODO(jessegross): LoRA loading
801
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
802
		panic("loras are not yet implemented")
803
804
	}

Jesse Gross's avatar
Jesse Gross committed
805
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
806
807
808
	if err != nil {
		panic(err)
	}
809

Jesse Gross's avatar
Jesse Gross committed
810
811
812
813
814
815
816
817
818
	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))

819
820
821
822
	s.status = ServerStatusReady
	s.ready.Done()
}

823
824
825
826
827
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")
828
829
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
830
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
831
832
833
	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")
834
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
835
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
836
837
	_ = 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")
838
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
839
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
840

841
	var lpaths multiLPath
842
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
843

844
845
846
847
848
849
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
	}
	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
867
	slog.Info("starting ollama engine")
868
869
870
871
872
873

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

Jesse Gross's avatar
Jesse Gross committed
874
875
876
877
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

878
	var tensorSplitFloats []float32
879
	if *tensorSplit != "" {
880
881
882
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
883
			f, _ := strconv.ParseFloat(s, 32)
884
			tensorSplitFloats[i] = float32(f)
885
		}
886
887
888
	}

	params := ml.BackendParams{
889
890
891
892
893
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
894
	}
895
896

	server.ready.Add(1)
897
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
898
899
900
901

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

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

904
905
906
907
908
909
	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)
910
		return err
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
	}
	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)
926
		return err
927
928
	}

929
	return nil
930
}