runner.go 22.4 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

	// vocab is a llama.cpp vocab required for gammar-based
	// constrained generation (json mode, structured outputs)
	// TODO: this is temporary until Ollama sampling supports
	// constrained generation
	vocab *sample.Vocab
263
264
265
266
267
268
269
270
271
272
273
274
}

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

func flushPending(seq *Sequence) bool {
275
276
277
278
279
280
281
282
283
284
285
	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]
286
287
	}

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

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
298
299
300
301
302
303
304
305
306
307
}

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
308
	seq.ctx.Close()
309
	s.seqs[seqIndex] = nil
310
	s.seqsSem.Release(1)
311
312
313
314
315
316
317
318
319
320
}

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

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

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

336
	var options input.Options
337

338
	for i, seq := range s.seqs {
339
340
341
342
343
		if seq == nil {
			continue
		}

		// if past the num predict limit
344
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
345
			s.removeSequence(i, "limit")
346
347
348
			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 = []input.Input{}
Jesse Gross's avatar
Jesse Gross committed
352
353
		}

354
		for j, inp := 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
				}
			}

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

370
371
372
			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})
373
374
			}

Jesse Gross's avatar
Jesse Gross committed
375
376
377
378
			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)
379
			if j+1 == len(seq.inputs) {
Jesse Gross's avatar
Jesse Gross committed
380
381
				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
			}
382
			seq.pendingInputs = append(seq.pendingInputs, inp)
383
		}
384
385

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
386
387
	}

Jesse Gross's avatar
Jesse Gross committed
388
	if len(options.Inputs) == 0 {
389
		return nil
390
391
	}

Jesse Gross's avatar
Jesse Gross committed
392
393
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()
394

Jesse Gross's avatar
Jesse Gross committed
395
	modelOutput, err := model.Forward(ctx, s.model, options)
396
	if err != nil {
397
		return fmt.Errorf("failed to decode batch: %w", err)
398
399
	}

400
	logits := modelOutput.Floats()
401

402
403
404
405
406
	for i, seq := range s.seqs {
		if seq == nil {
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
407
		// After calling Forward, pending inputs are now in the cache
408
409
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
410
			seq.pendingInputs = []input.Input{}
411
412
		}

413
414
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
415
416
417
			if !s.cache.enabled {
				return errors.New("caching disabled but unable to fit entire input in a batch")
			}
418
419
420
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
421
422
		seq.numPredicted++
		if seq.numPredicted == 1 {
423
424
425
426
427
			seq.startGenerationTime = time.Now()
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Jesse Gross's avatar
Jesse Gross committed
428
			// TODO(jessegross): Embedding support
429
430
431
			slog.Warn("generation of embedding outputs not yet supported")
			s.removeSequence(i, "")
			continue
432
433
434
		}

		// sample a token
435
436
437
		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
438
		if err != nil {
439
			return fmt.Errorf("failed to sample token: %w", err)
Jesse Gross's avatar
Jesse Gross committed
440
		}
441
442

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

457
		seq.inputs = []input.Input{{Token: token}}
458
459
460
461

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

Jesse Gross's avatar
Jesse Gross committed
462
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
463
464
465
466
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
467
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
			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]
483
484
485
486
487

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

Jesse Gross's avatar
Jesse Gross committed
488
		if common.ContainsStopSuffix(sequence, seq.stop) {
489
490
491
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
492
		if common.IncompleteUnicode(sequence) {
493
494
495
496
497
498
499
			continue
		}

		if !flushPending(seq) {
			s.removeSequence(i, "connection")
		}
	}
500
501

	return nil
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
}

// 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 {
529
530
531
	Data          []byte `json:"data"`
	ID            int    `json:"id"`
	AspectRatioID int    `json:"aspect_ratio_id"`
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
577
578
579
580
581
582
}

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
	}

583
584
585
586
587
588
589
590
591
592
	var grammar *sample.Grammar
	var err error
	if req.Grammar != "" {
		grammar, err = sample.NewGrammar(s.vocab, req.Grammar)
		if err != nil {
			http.Error(w, "failed to load model vocabulary required for format", http.StatusInternalServerError)
			return
		}
	}

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

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

614
	// Ensure there is a place to put the sequence, released when removed from s.seqs
615
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
616
617
618
619
620
		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)
		}
621
622
623
		return
	}

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

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

643
644
645
646
647
	if !found {
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
	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
672
						PredictedN:  seq.numPredicted,
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
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
						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"`
}

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

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

739
740
func (s *Server) loadModel(
	mpath string,
741
	params ml.BackendParams,
742
	lpath multiLPath,
Jesse Gross's avatar
Jesse Gross committed
743
	parallel int,
744
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
745
	kvSize int,
746
747
	multiUserCache bool,
) {
748
	var err error
749
	s.model, err = model.New(mpath, params)
750
751
752
	if err != nil {
		panic(err)
	}
753

754
755
	s.vocab = sample.NewVocab(mpath)

Jesse Gross's avatar
Jesse Gross committed
756
	// TODO(jessegross): LoRA loading
757
	if lpath.String() != "" {
Jesse Gross's avatar
Jesse Gross committed
758
		panic("loras are not yet implemented")
759
760
	}

Jesse Gross's avatar
Jesse Gross committed
761
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
762
763
764
	if err != nil {
		panic(err)
	}
765

Jesse Gross's avatar
Jesse Gross committed
766
767
768
769
770
771
772
773
774
	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))

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

779
780
781
782
783
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")
784
785
	numGPULayers := fs.Int("n-gpu-layers", 0, "Number of layers to offload to GPU")
	mainGPU := fs.Int("main-gpu", 0, "Main GPU")
786
	flashAttention := fs.Bool("flash-attn", false, "Enable flash attention")
787
788
789
	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")
790
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
791
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
Jesse Gross's avatar
Jesse Gross committed
792
793
	_ = 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")
794
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
795
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")
796

797
	var lpaths multiLPath
798
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")
799

800
801
802
803
804
805
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
	}
	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
823
	slog.Info("starting ollama engine")
824
825
826
827
828
829

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

Jesse Gross's avatar
Jesse Gross committed
830
831
832
833
	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

834
	var tensorSplitFloats []float32
835
	if *tensorSplit != "" {
836
837
838
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
839
			f, _ := strconv.ParseFloat(s, 32)
840
			tensorSplitFloats[i] = float32(f)
841
		}
842
843
844
	}

	params := ml.BackendParams{
845
846
847
848
849
		NumThreads:     *threads,
		NumGPULayers:   *numGPULayers,
		MainGPU:        *mainGPU,
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
850
	}
851
852

	server.ready.Add(1)
853
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
854
855
856
857

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

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

860
861
862
863
864
865
	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)
866
		return err
867
868
869
870
	}
	defer listener.Close()

	mux := http.NewServeMux()
871
872
873
874
875
876
877
	// TODO: support embeddings
	mux.HandleFunc("POST /embedding", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
	})

	mux.HandleFunc("POST /completion", server.completion)
	mux.HandleFunc("GET /health", server.health)
878
879
880
881
882
883
884
885

	httpServer := http.Server{
		Handler: mux,
	}

	log.Println("Server listening on", addr)
	if err := httpServer.Serve(listener); err != nil {
		log.Fatal("server error:", err)
886
		return err
887
888
	}

889
	return nil
890
}