runner.go 37.8 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package ollamarunner
2
3

import (
4
	"bytes"
5
6
7
8
9
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
10
	"hash/maphash"
11
	"image"
12
13
14
15
16
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
Jesse Gross's avatar
Jesse Gross committed
17
	"reflect"
18
19
20
21
22
23
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
24
	"unicode/utf8"
25

26
	"golang.org/x/image/bmp"
27
28
	"golang.org/x/sync/semaphore"

29
	"github.com/ollama/ollama/api"
30
	"github.com/ollama/ollama/envconfig"
31
	"github.com/ollama/ollama/fs/ggml"
32
	"github.com/ollama/ollama/llm"
33
	"github.com/ollama/ollama/logutil"
34
	"github.com/ollama/ollama/ml"
Michael Yang's avatar
Michael Yang committed
35
	"github.com/ollama/ollama/ml/nn/pooling"
Jesse Gross's avatar
Jesse Gross committed
36
	"github.com/ollama/ollama/model"
37
	"github.com/ollama/ollama/model/input"
Jesse Gross's avatar
Jesse Gross committed
38
39
40
41
	"github.com/ollama/ollama/runner/common"
	"github.com/ollama/ollama/sample"

	_ "github.com/ollama/ollama/model/models"
42
43
44
)

type Sequence struct {
45
	// ctxs are used for allocating tensors that last the lifetime of the sequence, such as
46
	// multimodal embeddings
47
	ctxs []ml.Context
48

49
50
51
	// mmStore holds multimodal embeddings to mange memory and enable splitting across batches
	mmStore multimodalStore

52
53
54
55
	// batch index
	iBatch int

	// prompt inputs left to evaluate
56
	inputs []*input.Input
57

Jesse Gross's avatar
Jesse Gross committed
58
	// inputs that have been added to a batch but not yet submitted to Forward
59
	pendingInputs []*input.Input
60

61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
	// 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

76
77
	// sampler with transforms to run on generated logits
	sampler sample.Sampler
78
79
80
81
82
83
84
85

	// 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
86
	numKeep int32
87
88
89
90

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

91
92
93
	// shift if context window is exceeded
	shift bool

94
	doneReason llm.DoneReason
95
96

	// Metrics
97
98
99
100
101
	startedAt, lastUpdatedAt time.Time
	processingDuration       time.Duration
	samplingDuration         time.Duration
	numPredicted             int
	numPromptInputs          int
102
103
104
}

type NewSequenceParams struct {
Jesse Gross's avatar
Jesse Gross committed
105
106
107
	numPredict int
	stop       []string
	numKeep    int32
108
	sampler    sample.Sampler
Jesse Gross's avatar
Jesse Gross committed
109
	embedding  bool
110
111
	shift      bool
	truncate   bool
112
113
}

114
115
var errorInputTooLong = errors.New("the input length exceeds the context length")

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

119
	inputs, ctxs, mmStore, err := s.inputs(prompt, images)
120
121
122
123
124
125
126
	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
127
		params.numKeep = int32(len(inputs))
128
129
	}

130
131
132
	// 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
133
134
	if int32(len(inputs)) > s.cache.numCtx {
		discard := int32(len(inputs)) - s.cache.numCtx
135
136
137
138
139

		if !params.truncate {
			return nil, errorInputTooLong
		}

140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
		promptStart := params.numKeep + discard

		// If we need to truncate in the middle of a unbreakable batch, remove the entire batch
		sameBatch := 0
		for i, inp := range inputs {
			if sameBatch > 0 {
				sameBatch--

				if promptStart == int32(i) {
					promptStart++
				}
			} else if promptStart == int32(i) {
				break
			}

			if inp.SameBatch != 0 {
				if int32(i) < params.numKeep {
					return nil, fmt.Errorf("SameBatch may not be specified within numKeep (index: %v numKeep: %v SameBatch: %v)", i, params.numKeep, inp.SameBatch)
				}

				sameBatch = inp.SameBatch
			}
		}

		if promptStart >= int32(len(inputs)) {
			return nil, errors.New("entire prompt removed by truncation")
		}

168
		newInputs := inputs[:params.numKeep]
169
		newInputs = append(newInputs, inputs[promptStart:]...)
170
171

		slog.Warn("truncating input prompt", "limit", s.cache.numCtx, "prompt", len(inputs), "keep", params.numKeep, "new", len(newInputs))
172
		inputs = newInputs
173
174
	}

Jesse Gross's avatar
Jesse Gross committed
175
	// TODO(jessegross): Ingest cached history for grammar
176
177

	return &Sequence{
178
179
180
181
182
183
184
185
186
187
188
189
190
		ctxs:             ctxs,
		mmStore:          mmStore,
		inputs:           inputs,
		numPromptInputs:  len(inputs),
		numPredict:       params.numPredict,
		pendingResponses: make([]string, 0),
		responses:        make(chan string, 100),
		quit:             make(chan bool, 1),
		embedding:        make(chan []float32, 1),
		sampler:          params.sampler,
		embeddingOnly:    params.embedding,
		stop:             params.stop,
		numKeep:          params.numKeep,
191
		shift:            params.shift,
192
193
194
195
196
	}, 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
197
// decoding images
198
199
func (s *Server) inputs(prompt string, images []llm.ImageData) ([]*input.Input, []ml.Context, multimodalStore, error) {
	var inputs []*input.Input
200
	var ctxs []ml.Context
201
	var mmStore multimodalStore
202

203
204
205
	var parts []string
	var matches [][]string

206
	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
207

208
209
210
211
	if visionModel {
		re := regexp.MustCompile(`\[img-(\d+)\]`)
		parts = re.Split(prompt, -1)
		matches = re.FindAllStringSubmatch(prompt, -1)
212
		mmStore = newMultimodalStore()
213
214
215
216
	} else {
		parts = []string{prompt}
	}

217
218
	for i, part := range parts {
		// text - tokenize
219
		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
220
		if err != nil {
221
			return nil, nil, nil, err
222
		}
223

224
		for _, t := range tokens {
225
			inputs = append(inputs, &input.Input{Token: t})
226
227
		}

Jesse Gross's avatar
Jesse Gross committed
228
		// image - decode and store
229
230
231
232
233
234
235
236
237
238
239
240
		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 {
241
				return nil, nil, nil, fmt.Errorf("invalid image index: %d", n)
242
243
			}

244
			ctx := s.model.Backend().NewContext()
245
246
			runtime.SetFinalizer(ctx, func(c ml.Context) { c.Close() })
			ctxs = append(ctxs, ctx)
247
			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
Jesse Gross's avatar
Jesse Gross committed
248
			if err != nil {
249
				return nil, nil, nil, err
Jesse Gross's avatar
Jesse Gross committed
250
251
			}

252
253
254
255
			s.multimodalHash.Reset()
			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
			imageHash := s.multimodalHash.Sum64()

256
257
			mmStore.addMultimodal(imageEmbeddings)

258
			inputs = append(inputs, &input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
259
260
261
		}
	}

262
	if visionModel {
263
		var err error
264
		inputs, err = multimodalProcessor.PostTokenize(inputs)
265
		if err != nil {
266
			return nil, nil, nil, err
267
268
269
		}
	}

270
	return inputs, ctxs, mmStore, nil
271
272
}

273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
type batchState struct {
	// id provides a counter for trace logging batches
	id int

	// ctx holds the backend context used for this batch
	ctx ml.Context

	// modelOutput holds the outputs from this batch
	modelOutput ml.Tensor

	// batchInputs holds the input token pointers which may start as
	// placeholders later filled in before calling ctx.Compute
	batchInputs []*input.Input

	// batch contains the inputs for a model forward pass
	batch input.Batch

	// full set of seqs at the time this batch was initiated
	seqs []*Sequence

	// Signaled when this batches inputs are ready and compute can proceed
	inputsReadyCh chan struct{}

	// Signaling when Compute is about to begin on this batch, and
	// seqs have been updated to prepare for the next batch
	computeStartedCh chan struct{}

	// Signaled when this batches outputs are complete and the next batch can proceed
	outputsReadyCh chan struct{}
}

304
type Server struct {
Jesse Gross's avatar
Jesse Gross committed
305
306
307
308
309
310
311
312
313
314
	// 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

	// lastLoad is the load request from the previous load attempt. Used to
	// detect if we can reuse an existing memory allocation.
	lastLoad llm.LoadRequest

315
316
317
318
319
	// 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
320
	model model.Model
321

322
	// status for external health reporting - loading, ready to serve, etc.
323
	status llm.ServerStatus
324
325
326
327
328
329
330
331

	// 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)
332
	// TODO (jmorganca): make this n_batch
333
334
	batchSize int

335
336
337
	// Simple counter used only for trace logging batches
	batchID int

338
339
340
341
342
343
344
345
	// 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
346
347
	seqs []*Sequence

348
349
350
351
	// seqs can have a maximum of parallel entries, which
	// is enfoced by seqSem
	seqsSem *semaphore.Weighted

352
353
354
	// KV cache
	cache *InputCache

355
356
357
	// next sequence for prompt processing to avoid starvation
	nextSeq int

358
359
360
	// multimodalHash generates hashes for comparing equality
	// of non-text data
	multimodalHash maphash.Hash
361
362
363
364
365
366
367
368
369
370
371
372
}

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

func flushPending(seq *Sequence) bool {
373
374
375
376
377
378
379
380
381
382
383
	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]
384
385
	}

386
387
388
389
390
391
392
393
394
395
	if len(joined) == 0 {
		return true
	}

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
396
397
}

398
func (s *Server) removeSequence(seqIndex int, reason llm.DoneReason) {
399
400
401
402
403
404
405
406
	seq := s.seqs[seqIndex]

	flushPending(seq)
	seq.doneReason = reason
	close(seq.responses)
	close(seq.embedding)
	seq.cache.InUse = false
	s.seqs[seqIndex] = nil
407
	s.seqsSem.Release(1)
408
409
}

410
411
// track batch state between forwardBatch, computeBatch and predictForwardBatch

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

Michael Yang's avatar
Michael Yang committed
415
	supportsAsync := pooling.Type(s.model.Backend().Config().Uint("pooling_type")) == pooling.TypeNone
Michael Yang's avatar
Michael Yang committed
416

417
	var previousBatch batchState
418
419
420
421
422
	for {
		select {
		case <-ctx.Done():
			return
		default:
423
			var err error
424
			nextBatch, err := s.forwardBatch(previousBatch)
425
426
427
			if err != nil {
				panic(err)
			}
Michael Yang's avatar
Michael Yang committed
428
429

			if supportsAsync {
430
				go s.computeBatch(nextBatch)
Michael Yang's avatar
Michael Yang committed
431
			} else {
432
				s.computeBatch(nextBatch)
Michael Yang's avatar
Michael Yang committed
433
			}
434
435

			previousBatch = nextBatch
436
437
438
439
		}
	}
}

440
441
442
443
444
445
// forwardBatch will calculate a batch.
func (s *Server) forwardBatch(pendingBatch batchState) (nextBatch batchState, err error) {
	// If we have a pending batch still processing, wait until Compute has started
	// before setting up the next batch so the seqs inputs are ready to receive their
	// token values and we get the correct input pointers for the batchInputs
	if pendingBatch.ctx != nil {
Michael Yang's avatar
Michael Yang committed
446
		logutil.Trace("forwardBatch waiting for compute to start", "pendingBatch.id", pendingBatch.id)
447
		<-pendingBatch.computeStartedCh
Michael Yang's avatar
Michael Yang committed
448
		logutil.Trace("forwardBatch compute started, setting up next batch", "pendingBatch.id", pendingBatch.id, "id", s.batchID)
449
450
		nextBatch.inputsReadyCh = pendingBatch.outputsReadyCh // Chain the ouputs from the pending batch to the next inputs batch
	} else {
Michael Yang's avatar
Michael Yang committed
451
		logutil.Trace("forwardBatch no pending batch detected", "batchID", s.batchID)
452
453
454
455
456
		// No pendingBatch, so the inputs will be ready in the seqs immediately
		nextBatch.inputsReadyCh = make(chan struct{}, 1)
		nextBatch.inputsReadyCh <- struct{}{}
	}

457
458
459
460
461
462
	s.mu.Lock()
	for s.allNil() {
		s.cond.Wait() // Wait until an item is added
	}
	defer s.mu.Unlock()

463
464
465
466
467
468
469
470
471
472
473
	nextBatch.ctx = s.model.Backend().NewContext()
	defer func() {
		if err != nil {
			nextBatch.ctx.Close()
			nextBatch.ctx = nil
		}
	}()
	nextBatch.id = s.batchID
	nextBatch.seqs = append([]*Sequence{}, s.seqs...)
	nextBatch.computeStartedCh = make(chan struct{}, 1)
	nextBatch.outputsReadyCh = make(chan struct{}, 1)
474

475
476
	// Prepare the seqs and batch, but defer the input token values as we may not be ready yet
	var batchInputs []*input.Input
477
	var batchOutputs []int32
Jesse Gross's avatar
Jesse Gross committed
478
	var batch input.Batch
479

480
481
482
483
484
	resumeSeq := -1
	seqIdx := s.nextSeq - 1
	for range s.seqs {
		seqIdx = (seqIdx + 1) % len(s.seqs)
		seq := s.seqs[seqIdx]
485
486
487
488
489
		if seq == nil {
			continue
		}

		// if past the num predict limit
490
		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
491
			s.removeSequence(seqIdx, llm.DoneReasonLength)
492
			nextBatch.seqs[seqIdx] = nil
493
494
495
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
496
497
		if !s.cache.enabled {
			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
498
			seq.cache.Inputs = []*input.Input{}
Jesse Gross's avatar
Jesse Gross committed
499
500
		}

501
502
		batchSize := s.batchSize

503
		for i, inp := range seq.inputs {
504
505
			// If we are required to put following inputs into a single batch then extend the
			// batch size. Since we are only extending the size the minimum amount possible, this
506
			// will cause a break if we have existing inputs.
507
508
509
510
511
			minBatch := 1 + inp.SameBatch
			if minBatch > batchSize {
				batchSize = minBatch
			}

512
513
514
515
516
517
518
519
			// Stop if the required batch would put us over the total batch size (including tokens
			// added by other sequences). If we haven't been able to add anything yet then pick up
			// here again for the next batch to avoid starvation, though we can opportunistically
			// check if other sequences can still squeeze something in.
			if len(batchInputs)+minBatch > batchSize {
				if len(seq.pendingInputs) == 0 && resumeSeq == -1 {
					resumeSeq = seqIdx
				}
520
521
				break
			}
Jesse Gross's avatar
Jesse Gross committed
522

523
524
525
526
527
528
529
530
			// If the sum of our working set (already processed tokens, tokens we added to this
			// batch, required following tokens) exceeds the context size, then trigger a shift
			// now so we don't have to do one later when we can't break the batch.
			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+minBatch) > s.cache.numCtx {
				if len(seq.pendingInputs) != 0 {
					break
				}

531
532
533
534
535
536
				if !seq.shift {
					s.removeSequence(seqIdx, llm.DoneReasonLength)
					nextBatch.seqs[seqIdx] = nil
					break
				}

537
				err = s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
538
				if err != nil {
539
540
541
542
543
					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...)
						// Skip this sequence but continue processing the rest
544
545
						nextBatch.seqs[seqIdx] = nil // clear this sequence for this batch
						err = nil
546
547
						continue
					} else {
548
						return
549
					}
550
551
552
				}
			}

553
			batchInputs = append(batchInputs, seq.inputs[i])
554
			if inp.Multimodal != nil {
555
556
				var mm []input.Multimodal
				mm, err = seq.mmStore.getMultimodal(s.model.Backend(), nextBatch.ctx, inp.Multimodal, false)
557
				if err != nil {
558
					return
559
560
				}
				batch.Multimodal = append(batch.Multimodal, input.MultimodalIndex{Index: len(batchInputs) - 1, Multimodal: mm})
561
562
			}

Jesse Gross's avatar
Jesse Gross committed
563
564
			batch.Positions = append(batch.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
			batch.Sequences = append(batch.Sequences, seq.cache.Id)
Jesse Gross's avatar
Jesse Gross committed
565

566
567
568
			seq.iBatch = len(batchOutputs)
			if i+1 == len(seq.inputs) || seq.embeddingOnly {
				batchOutputs = append(batchOutputs, int32(len(batchInputs)-1))
Jesse Gross's avatar
Jesse Gross committed
569
			}
Michael Yang's avatar
Michael Yang committed
570
			logutil.Trace("forwardBatch iBatch", "batchID", s.batchID, "seqIdx", seqIdx, "seq.iBatch", seq.iBatch, "i+1", i+1, "len(seq.inputs)", len(seq.inputs))
571
			seq.pendingInputs = append(seq.pendingInputs, inp)
572
		}
573
574

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
575
576
	}

577
578
579
580
581
582
583
	startedAt := time.Now()
	for i := range nextBatch.seqs {
		if nextBatch.seqs[i] != nil && nextBatch.seqs[i].startedAt.IsZero() {
			nextBatch.seqs[i].startedAt = startedAt
		}
	}

584
585
586
587
588
589
	if resumeSeq != -1 {
		s.nextSeq = resumeSeq
	} else {
		s.nextSeq = seqIdx + 1
	}

590
	if len(batchInputs) == 0 {
Michael Yang's avatar
Michael Yang committed
591
		logutil.Trace("forwardBatch no batchInputs, going idle", "batchID", s.batchID)
592
593
594
		nextBatch.ctx.Close()
		nextBatch.ctx = nil
		return
595
	}
596
	s.batchID++
597

598
599
	// Actual batchInputs values will be injected into the batch.Inputs tensor before calling Compute
	batch.Inputs = nextBatch.ctx.Input().Empty(ml.DTypeI32, len(batchInputs))
Michael Yang's avatar
Michael Yang committed
600
	batch.Outputs = nextBatch.ctx.Input().FromInts(batchOutputs, len(batchOutputs))
601
	nextBatch.modelOutput, err = model.Forward(nextBatch.ctx, s.model, batch)
602
	if err != nil {
603
604
		err = fmt.Errorf("failed to build graph: %w", err)
		return
605
	}
606
607
	nextBatch.batchInputs = batchInputs
	nextBatch.batch = batch
608

609
610
611
612
613
614
615
616
617
618
619
620
	return
}

// Async processing of the next batch
func (s *Server) computeBatch(activeBatch batchState) {
	if activeBatch.ctx == nil {
		// Nothing to compute
		return
	}
	defer activeBatch.ctx.Close()

	// Wait until inputs are ready
Michael Yang's avatar
Michael Yang committed
621
	logutil.Trace("computeBatch: waiting for inputs to be ready", "batchID", activeBatch.id)
622
	<-activeBatch.inputsReadyCh
Michael Yang's avatar
Michael Yang committed
623
	logutil.Trace("computeBatch: inputs are ready", "batchID", activeBatch.id)
624

625
626
627
	// Once we complete, signal the next batch of inputs are ready
	// This will unblock the next computeBatch, or forwardBatch if new seqs come in
	defer func() {
Michael Yang's avatar
Michael Yang committed
628
		logutil.Trace("computeBatch: outputs are ready", "batchID", activeBatch.id)
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
		activeBatch.outputsReadyCh <- struct{}{}
	}()

	s.mu.Lock()

	// Gather the actual input token values now that they're ready
	batchInputs := make([]int32, len(activeBatch.batchInputs))
	for i := range batchInputs {
		batchInputs[i] = activeBatch.batchInputs[i].Token
	}

	// Now we run part of the decoding algorithm to adjust the seq.inputs with placeholder tokens
	// so that forwardBatch can build a batchInputs set which will eventually contain the actual
	// decoded tokens.
	nextBatchTokens := make([]*input.Input, len(s.seqs))
	iBatches := make([]int, len(s.seqs)) // Record the iBatch values before releasing the lock
645
	for i, seq := range s.seqs {
646
		iBatches[i] = -1
647
648
649
		if seq == nil {
			continue
		}
650
651
652
653
		// Skip over any newly added or skipped sequences
		if activeBatch.seqs[i] == nil {
			continue
		}
654

655
656
657
		// Detect if the sequence we're processing has already been completed and replaced
		// with a new sequence
		if seq != activeBatch.seqs[i] {
Michael Yang's avatar
Michael Yang committed
658
			logutil.Trace("computeBatch: sequence replaced, discarding its results", "batchID", activeBatch.id, "seqIdx", i)
659
660
661
662
663
664
665
666
667
668
669
			continue
		}

		// Pending inputs will actually be in the cache after we call Compute.
		// However, we have already resolved any placeholder tokens.
		//
		// It's possible for incoming sequences to look at the values that we've
		// added to the cache here and start relying on them before we've done
		// the computation. This is OK as long as we ensure that this batch's
		// computation happens before any future batch's and we never fail
		// (unless we take down the whole runner).
670
671
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
672
			seq.pendingInputs = []*input.Input{}
673
674
		}

675
676
		// don't sample prompt processing
		if len(seq.inputs) != 0 {
Jesse Gross's avatar
Jesse Gross committed
677
			if !s.cache.enabled {
Michael Yang's avatar
Michael Yang committed
678
				panic("caching disabled but unable to fit entire input in a batch")
Jesse Gross's avatar
Jesse Gross committed
679
			}
680
681
682
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
683
		seq.numPredicted++
684
685
686
687
688
689
690
691
692
		nextToken := &input.Input{Token: 0} // placeholder we'll fill in after Compute/Floats
		seq.inputs = []*input.Input{nextToken}
		nextBatchTokens[i] = nextToken
		iBatches[i] = seq.iBatch
	}

	// At this point the seqs are ready for forwardBatch to move forward so unblock
	s.mu.Unlock()

Michael Yang's avatar
Michael Yang committed
693
	activeBatch.batch.Inputs.FromInts(batchInputs)
694
695
	activeBatch.ctx.ComputeWithNotify(
		func() {
Michael Yang's avatar
Michael Yang committed
696
			logutil.Trace("computeBatch: signaling computeStartedCh", "batchID", activeBatch.id)
697
698
699
			activeBatch.computeStartedCh <- struct{}{}
		},
		activeBatch.modelOutput)
Michael Yang's avatar
Michael Yang committed
700
701

	outputs := activeBatch.modelOutput.Floats()
702
	t := time.Now()
703

Michael Yang's avatar
Michael Yang committed
704
	logutil.Trace("computeBatch: logits ready", "batchID", activeBatch.id)
705
706
707
708

	s.mu.Lock()
	defer s.mu.Unlock()

Michael Yang's avatar
Michael Yang committed
709
	logutil.Trace("computeBatch: decoding", "batchID", activeBatch.id)
710
711
712
713
714
	for i, seq := range s.seqs {
		if seq == nil || nextBatchTokens[i] == nil {
			continue
		}

715
		seq.lastUpdatedAt = t
Jesse Gross's avatar
Jesse Gross committed
716
		if seq.numPredicted == 1 {
717
718
			seq.processingDuration = seq.lastUpdatedAt.Sub(seq.startedAt)
			seq.startedAt = seq.lastUpdatedAt
719
720
721
722
		}

		// if done processing the prompt, generate an embedding and return
		if seq.embeddingOnly {
Michael Yang's avatar
Michael Yang committed
723
			seq.embedding <- outputs
724
			s.removeSequence(i, llm.DoneReasonStop)
725
			continue
726
727
728
		}

		// sample a token
729
730
		vocabSize := len(outputs) / activeBatch.batch.Outputs.Dim(0)
		logutil.Trace("computeBatch: vocab details", "batchID", activeBatch.id, "seqIdx", i, "len(logits)", len(outputs), "len(activeBatch.batch.Outputs)", activeBatch.batch.Outputs.Dim(0), "vocabSize", vocabSize, "iBatches", iBatches)
Michael Yang's avatar
Michael Yang committed
731
		token, err := seq.sampler.Sample(outputs[iBatches[i]*vocabSize : (iBatches[i]+1)*vocabSize])
Jesse Gross's avatar
Jesse Gross committed
732
		if err != nil {
Michael Yang's avatar
Michael Yang committed
733
			panic("failed to sample token")
Jesse Gross's avatar
Jesse Gross committed
734
		}
735

736
737
		nextBatchTokens[i].Token = token

738
		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
739
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
740
741
742
			// TODO (jmorganca): we should send this back
			// as it's important for the /api/generate context
			// seq.responses <- piece
Michael Yang's avatar
Michael Yang committed
743
			logutil.Trace("computeBatch: EOS", "batchID", activeBatch.id, "seqIdx", i)
744
			s.removeSequence(i, llm.DoneReasonStop)
745
746
747
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
748
749
		piece, err := s.model.(model.TextProcessor).Decode([]int32{token})
		if err != nil {
Michael Yang's avatar
Michael Yang committed
750
			panic("failed to decode token")
Jesse Gross's avatar
Jesse Gross committed
751
752
		}

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

Jesse Gross's avatar
Jesse Gross committed
756
		if ok, stop := common.FindStop(sequence, seq.stop); ok {
757
758
759
760
			slog.Debug("hit stop token", "pending", seq.pendingResponses, "stop", stop)

			var tokenTruncated bool
			origLen := len(seq.pendingResponses)
Jesse Gross's avatar
Jesse Gross committed
761
			seq.pendingResponses, tokenTruncated = common.TruncateStop(seq.pendingResponses, stop)
762
763
764
765
766
767
768
769
770
771
772
773
774
775
			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--
			}
776

777
			seq.cache.Inputs = seq.cache.Inputs[:tokenLen]
778

779
			s.removeSequence(i, llm.DoneReasonStop)
780
781
782
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
783
		if common.ContainsStopSuffix(sequence, seq.stop) {
784
785
786
			continue
		}

Jesse Gross's avatar
Jesse Gross committed
787
		if common.IncompleteUnicode(sequence) {
788
789
790
791
			continue
		}

		if !flushPending(seq) {
792
			s.removeSequence(i, llm.DoneReasonConnectionClosed)
793
794
		}
	}
795
796
797
798
799
800
801

	samplingDuration := time.Since(t)
	for i, seq := range s.seqs {
		if seq != nil && nextBatchTokens[i] != nil {
			s.seqs[i].samplingDuration += samplingDuration
		}
	}
802
803
804
}

func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
805
	var req llm.CompletionRequest
806
807
808
809
810
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

811
812
813
814
815
	if req.Options == nil {
		opts := api.DefaultOptions()
		req.Options = &opts
	}

816
817
818
819
820
821
822
823
824
825
	// 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
	}

826
	var grammar *sample.GrammarSampler
827
828
	var err error
	if req.Grammar != "" {
829
		grammar, err = sample.NewGrammarSampler(s.model.(model.TextProcessor), req.Grammar)
830
831
832
833
		if err != nil {
			http.Error(w, "failed to load model vocabulary required for format", http.StatusInternalServerError)
			return
		}
834
		defer grammar.Free()
835
836
	}

837
	sampler := sample.NewSampler(
838
839
840
841
842
		req.Options.Temperature,
		req.Options.TopK,
		req.Options.TopP,
		req.Options.MinP,
		req.Options.Seed,
843
		grammar,
844
845
	)

846
	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
847
848
849
		numPredict: req.Options.NumPredict,
		stop:       req.Options.Stop,
		numKeep:    int32(req.Options.NumKeep),
850
		sampler:    sampler,
Jesse Gross's avatar
Jesse Gross committed
851
		embedding:  false,
852
853
		shift:      req.Shift,
		truncate:   req.Truncate,
854
855
	})
	if err != nil {
856
857
858
859
		if errors.Is(err, errorInputTooLong) {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
860
861
862
863
		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

864
	// Ensure there is a place to put the sequence, released when removed from s.seqs
865
	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
866
867
868
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting completion request due to client closing the connection")
		} else {
869
			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
870
		}
871
872
873
		return
	}

874
	s.mu.Lock()
875
	found := false
876
877
	for i, sq := range s.seqs {
		if sq == nil {
Michael Yang's avatar
Michael Yang committed
878
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
879
880
			if err != nil {
				s.mu.Unlock()
881
				s.seqsSem.Release(1)
882
883
884
				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}
885

886
887
			s.seqs[i] = seq
			s.cond.Signal()
888
			found = true
889
890
891
892
893
			break
		}
	}
	s.mu.Unlock()

894
	if !found {
895
		s.seqsSem.Release(1)
896
897
898
899
		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
		return
	}

900
901
902
903
904
905
906
	for {
		select {
		case <-r.Context().Done():
			close(seq.quit)
			return
		case content, ok := <-seq.responses:
			if ok {
907
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
908
					Content: content,
909
910
911
912
913
914
915
916
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
					close(seq.quit)
					return
				}

				flusher.Flush()
			} else {
917
918
				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
					Done:               true,
919
					DoneReason:         seq.doneReason,
920
					PromptEvalCount:    seq.numPromptInputs,
921
					PromptEvalDuration: seq.processingDuration,
922
					EvalCount:          seq.numPredicted,
923
					EvalDuration:       seq.lastUpdatedAt.Sub(seq.startedAt) - seq.samplingDuration,
924
925
926
927
928
929
930
931
932
933
				}); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
				}

				return
			}
		}
	}
}

Michael Yang's avatar
Michael Yang committed
934
func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
Michael Yang's avatar
Michael Yang committed
935
	if pooling.Type(s.model.Backend().Config().Uint("pooling_type")) == pooling.TypeNone {
Michael Yang's avatar
Michael Yang committed
936
937
938
939
940
941
942
943
944
945
946
		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
		return
	}

	var req llm.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")
947
948
	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{
		embedding: true,
949
950
951
952
953

		// TODO (jmorganca): this should be provided by the server via the
		// request options and truncated here in the runner, instead of relying on
		// the server's truncate logic
		truncate: true,
954
	})
Michael Yang's avatar
Michael Yang committed
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
	if err != nil {
		http.Error(w, fmt.Sprintf("failed to create new sequence: %v", err), http.StatusInternalServerError)
		return
	}

	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
		if errors.Is(err, context.Canceled) {
			slog.Info("aborting embedding request due to client closing the connection")
		} else {
			http.Error(w, fmt.Sprintf("failed to acquire semaphore: %v", err), http.StatusInternalServerError)
		}
		return
	}

	s.mu.Lock()
	found := false
	for i, sq := range s.seqs {
		if sq == nil {
			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
			if err != nil {
				s.mu.Unlock()
				s.seqsSem.Release(1)
				http.Error(w, fmt.Sprintf("failed to load cache: %v", err), http.StatusInternalServerError)
				return
			}

			s.seqs[i] = seq
			s.cond.Signal()
			found = true
			break
		}
	}
	s.mu.Unlock()

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

	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
996
		Embedding: <-seq.embedding,
Michael Yang's avatar
Michael Yang committed
997
998
999
1000
1001
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

1002
1003
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
1004
1005
	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
		Status:   s.status,
1006
1007
1008
1009
1010
1011
		Progress: s.progress,
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

1012
func (s *Server) reserveWorstCaseGraph() error {
1013
1014
1015
	ctx := s.model.Backend().NewContext()
	defer ctx.Close()

1016
	var err error
1017
1018
1019
1020
	inputs := make([]*input.Input, s.batchSize)
	for i := range inputs {
		inputs[i] = &input.Input{}
	}
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
	mmStore := newMultimodalStore()

	// Multimodal strategy:
	// - Encode a 2048x2048 image. This assumes that a single image of this
	//   size is sufficient to trigger the worst case. This is currently true
	//   because for existing models, only a single image fits in a batch.
	// - Add the embedding to a full batch of tokens - this is necessary because
	//   the model may be looking for non-image data, such as <image> tags.
	// - Run PostTokenize to execute any transformations between generated
	//   embeddings and what the forward pass expects.
	// - The result may now be larger than a batch (images may not fit in a
	//   single batch), so trim based on what will fit and must be grouped together.
	// - Fill out the rest of the space with text tokens.
	if multimodalProcessor, ok := s.model.(model.MultimodalProcessor); ok {
		mmCtx := s.model.Backend().NewContext()
		defer mmCtx.Close()

		img := image.NewGray(image.Rect(0, 0, 2048, 2048))
		var buf bytes.Buffer
		bmp.Encode(&buf, img)

		if inputs[0].Multimodal, err = multimodalProcessor.EncodeMultimodal(mmCtx, buf.Bytes()); err == nil {
			mmStore.addMultimodal(inputs[0].Multimodal)

			inputs, err = multimodalProcessor.PostTokenize(inputs)
			if err != nil {
				return err
			}

			for i, inp := range inputs {
				minBatch := 1 + inp.SameBatch
				if minBatch > s.batchSize {
					inputs = inputs[i:min(i+minBatch, len(inputs))]
					break
				} else if i+minBatch > s.batchSize {
					inputs = inputs[:i]
					break
				}
			}

			if len(inputs) < s.batchSize {
1062
				newInputs := make([]*input.Input, s.batchSize)
1063
				copy(newInputs, inputs)
1064
1065
1066
				for i := len(inputs); i < s.batchSize; i++ {
					newInputs[i] = &input.Input{}
				}
1067
1068
1069
1070
1071
				inputs = newInputs
			}
		}
	}

1072
1073
	var batch input.Batch

1074
	batchInputs := make([]int32, len(inputs))
1075
1076
	batch.Positions = make([]int32, len(inputs))
	batch.Sequences = make([]int, len(inputs))
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
	for i, inp := range inputs {
		batchInputs[i] = inp.Token
		if inp.Multimodal != nil {
			mm, err := mmStore.getMultimodal(s.model.Backend(), ctx, inp.Multimodal, true)
			if err != nil {
				return err
			}
			batch.Multimodal = append(batch.Multimodal, input.MultimodalIndex{Index: i, Multimodal: mm})
		}

1087
1088
1089
		batch.Positions[i] = int32(i)
	}

Michael Yang's avatar
Michael Yang committed
1090
	batch.Inputs = ctx.Input().FromInts(batchInputs, len(batchInputs))
1091
	batch.Outputs = ctx.Input().Empty(ml.DTypeI32, s.parallel)
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105

	cache := s.model.Config().Cache
	if cache != nil {
		err := cache.StartForward(ctx, batch, true)
		if err != nil {
			return err
		}
	}

	t, err := s.model.Forward(ctx, batch)
	if err != nil {
		return err
	}

1106
	ctx.Forward(t).Reserve()
1107
1108

	return nil
1109
}
1110

Jesse Gross's avatar
Jesse Gross committed
1111
1112
1113
// allocModel pre-allocates the maximum needed memory for a model
// based on the given parameters
func (s *Server) allocModel(
1114
	mpath string,
1115
	params ml.BackendParams,
Jesse Gross's avatar
Jesse Gross committed
1116
	loraPath []string,
Jesse Gross's avatar
Jesse Gross committed
1117
	parallel int,
1118
	kvCacheType string,
Jesse Gross's avatar
Jesse Gross committed
1119
	kvSize int,
1120
	multiUserCache bool,
Jesse Gross's avatar
Jesse Gross committed
1121
1122
1123
1124
1125
) (panicErr error) {
	// Convert memory allocation panics to errors
	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(error); ok {
1126
1127
1128
1129
1130
1131
				var noMem ml.ErrNoMem
				if errors.As(err, &noMem) {
					panicErr = noMem
				} else {
					panic(r)
				}
Jesse Gross's avatar
Jesse Gross committed
1132
1133
1134
1135
1136
1137
			} else {
				panic(r)
			}
		}
	}()

1138
	var err error
1139
	s.model, err = model.New(mpath, params)
1140
	if err != nil {
1141
		return err
1142
	}
1143

Jesse Gross's avatar
Jesse Gross committed
1144
	// TODO(jessegross): LoRA loading
Jesse Gross's avatar
Jesse Gross committed
1145
	if len(loraPath) > 0 {
1146
		return errors.New("loras are not yet implemented")
1147
1148
	}

1149
	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, s.batchSize, multiUserCache)
1150
	if err != nil {
1151
		return err
1152
	}
1153

Jesse Gross's avatar
Jesse Gross committed
1154
1155
1156
1157
1158
1159
1160
1161
1162
	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))

1163
1164
1165
	return s.reserveWorstCaseGraph()
}

Jesse Gross's avatar
Jesse Gross committed
1166
1167
1168
1169
1170
1171
1172
// closeModel frees all memory associated with a model
func (s *Server) closeModel() {
	s.cache.Close()
	s.cache = nil
	if s.model != nil {
		s.model.Backend().Close()
		s.model = nil
1173
	}
Jesse Gross's avatar
Jesse Gross committed
1174
}
1175

Jesse Gross's avatar
Jesse Gross committed
1176
1177
1178
1179
// loadModel loads the weights for a model. The memory must already
// have been allocated with allocModel
func (s *Server) loadModel() {
	err := s.model.Backend().Load(context.TODO(),
1180
1181
1182
1183
		func(progress float32) {
			s.progress = progress
		})
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
1184
		panic(fmt.Errorf("failed to load model: %v", err))
1185
1186
	}

1187
	s.status = llm.ServerStatusReady
1188
1189
1190
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
// 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)

	if req.Operation == llm.LoadOperationClose {
		s.closeModel()
		if err := json.NewEncoder(w).Encode(&llm.LoadResponse{}); err != nil {
			http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		}
		return
	}

	s.lastLoad.Operation = req.Operation
	loadModel := s.model == nil || !reflect.DeepEqual(req, s.lastLoad)

	s.lastLoad = req

	if loadModel {
		s.closeModel()

		params := ml.BackendParams{
			AllocMemory:    req.Operation != llm.LoadOperationFit,
			NumThreads:     req.NumThreads,
			GPULayers:      req.GPULayers,
			FlashAttention: req.FlashAttention,
		}

		s.batchSize = req.BatchSize

		err := s.allocModel(s.modelPath, params, req.LoraPath, req.Parallel, req.KvCacheType, req.KvSize, req.MultiUserCache)
		if err != nil {
			s.closeModel()

			var noMem ml.ErrNoMem
			if errors.As(err, &noMem) {
				resp := llm.LoadResponse{Success: false, Memory: noMem.BackendMemory}
				if err := json.NewEncoder(w).Encode(&resp); err != nil {
					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
				}

				return
			}

			http.Error(w, fmt.Sprintf("failed to initialize model: %v", err), http.StatusInternalServerError)
			return
		}
	}

	mem := s.model.Backend().BackendMemory()

	switch req.Operation {
	case llm.LoadOperationFit:
		// LoadOperationFit can't be used for anything else, so just close it
		s.closeModel()

	// LoadOperationAlloc should stay open for future operations

	case llm.LoadOperationCommit:
		s.status = llm.ServerStatusLoadingModel
		go s.loadModel()
	}

	resp := llm.LoadResponse{Success: true, Memory: mem}
	if err := json.NewEncoder(w).Encode(&resp); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		return
	}
}

1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
// info is the handler called by the Ollama server to report information
// about the GPU devices in use by this runner
func (s *Server) info(w http.ResponseWriter, r *http.Request) {
	s.loadMu.Lock()
	defer s.loadMu.Unlock()

	w.Header().Set("Content-Type", "application/json")

	m := s.model

	if m == nil {
		startLoad := time.Now()

		// Dummy load to get the backend wired up
		f, err := os.CreateTemp("", "*.bin")
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to initialize baackend: %v", err), http.StatusInternalServerError)
			return
		}
		defer f.Close()
		defer os.Remove(f.Name())

		if err := ggml.WriteGGUF(f, ggml.KV{
			"general.architecture": "llama",
			"tokenizer.ggml.model": "gpt2",
		}, nil); err != nil {
			http.Error(w, fmt.Sprintf("failed to initialize baackend: %v", err), http.StatusInternalServerError)
			return
		}

		m, err = model.New(f.Name(), ml.BackendParams{NumThreads: runtime.NumCPU(), AllocMemory: false, GPULayers: ml.GPULayersList{{}}})
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to initialize baackend: %v", err), http.StatusInternalServerError)
			return
		}
		slog.Debug("dummy model load took", "duration", time.Since(startLoad))
	}

	startDevices := time.Now()
	infos := m.Backend().BackendDevices()
	slog.Debug("gathering device infos took", "duration", time.Since(startDevices))
	if err := json.NewEncoder(w).Encode(&infos); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

1323
1324
1325
1326
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")
1327
	_ = fs.Bool("verbose", false, "verbose output (default: disabled)")
1328

1329
1330
1331
1332
1333
1334
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
1335
	}
1336
	slog.SetDefault(logutil.NewLogger(os.Stderr, envconfig.LogLevel()))
Jesse Gross's avatar
Jesse Gross committed
1337
	slog.Info("starting ollama engine")
1338

1339
1340
1341
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

Jesse Gross's avatar
Jesse Gross committed
1342
1343
1344
	server := &Server{
		modelPath: *mpath,
		status:    llm.ServerStatusLaunched,
1345
1346
	}

Jesse Gross's avatar
Jesse Gross committed
1347
1348
	server.cond = sync.NewCond(&server.mu)
	server.ready.Add(1)
1349
1350
1351
1352
1353
1354
1355

	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)
1356
		return err
1357
1358
1359
1360
	}
	defer listener.Close()

	mux := http.NewServeMux()
1361
	// TODO: support embeddings
1362
	mux.HandleFunc("GET /info", server.info)
Jesse Gross's avatar
Jesse Gross committed
1363
	mux.HandleFunc("POST /load", server.load)
Michael Yang's avatar
Michael Yang committed
1364
	mux.HandleFunc("POST /embedding", server.embeddings)
1365
1366
	mux.HandleFunc("POST /completion", server.completion)
	mux.HandleFunc("GET /health", server.health)
1367
1368
1369
1370
1371
1372
1373
1374

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

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

1378
	return nil
1379
}