runner.go 37.9 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
217
	} else {
		parts = []string{prompt}
	}

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

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

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

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

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

257
258
			mmStore.addMultimodal(imageEmbeddings)

259
			inputs = append(inputs, &input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
260
261
262
263
264
265
			postTokenize = true
		}
	}

	if visionModel && postTokenize {
		var err error
266
		inputs, err = multimodalProcessor.PostTokenize(inputs)
267
		if err != nil {
268
			return nil, nil, nil, err
269
270
271
		}
	}

272
	return inputs, ctxs, mmStore, nil
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
304
305
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{}
}

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

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

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

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

337
338
339
	// Simple counter used only for trace logging batches
	batchID int

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

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

354
355
356
	// KV cache
	cache *InputCache

357
358
359
	// next sequence for prompt processing to avoid starvation
	nextSeq int

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

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

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

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

	select {
	case seq.responses <- joined:
		return true
	case <-seq.quit:
		return false
	}
398
399
}

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

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

412
413
// track batch state between forwardBatch, computeBatch and predictForwardBatch

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

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

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

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

			previousBatch = nextBatch
438
439
440
441
		}
	}
}

442
443
444
445
446
447
// 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
448
		logutil.Trace("forwardBatch waiting for compute to start", "pendingBatch.id", pendingBatch.id)
449
		<-pendingBatch.computeStartedCh
Michael Yang's avatar
Michael Yang committed
450
		logutil.Trace("forwardBatch compute started, setting up next batch", "pendingBatch.id", pendingBatch.id, "id", s.batchID)
451
452
		nextBatch.inputsReadyCh = pendingBatch.outputsReadyCh // Chain the ouputs from the pending batch to the next inputs batch
	} else {
Michael Yang's avatar
Michael Yang committed
453
		logutil.Trace("forwardBatch no pending batch detected", "batchID", s.batchID)
454
455
456
457
458
		// No pendingBatch, so the inputs will be ready in the seqs immediately
		nextBatch.inputsReadyCh = make(chan struct{}, 1)
		nextBatch.inputsReadyCh <- struct{}{}
	}

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

465
466
467
468
469
470
471
472
473
474
475
	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)
476

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

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

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

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

503
504
		batchSize := s.batchSize

505
		for i, inp := range seq.inputs {
506
507
			// 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
508
			// will cause a break if we have existing inputs.
509
510
511
512
513
			minBatch := 1 + inp.SameBatch
			if minBatch > batchSize {
				batchSize = minBatch
			}

514
515
516
517
518
519
520
521
			// 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
				}
522
523
				break
			}
Jesse Gross's avatar
Jesse Gross committed
524

525
526
527
528
529
530
531
532
			// 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
				}

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

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

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

Jesse Gross's avatar
Jesse Gross committed
565
566
			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
567

568
569
570
			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
571
			}
Michael Yang's avatar
Michael Yang committed
572
			logutil.Trace("forwardBatch iBatch", "batchID", s.batchID, "seqIdx", seqIdx, "seq.iBatch", seq.iBatch, "i+1", i+1, "len(seq.inputs)", len(seq.inputs))
573
			seq.pendingInputs = append(seq.pendingInputs, inp)
574
		}
575
576

		seq.inputs = seq.inputs[len(seq.pendingInputs):]
577
578
	}

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

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

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

600
601
	// 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
602
	batch.Outputs = nextBatch.ctx.Input().FromInts(batchOutputs, len(batchOutputs))
603
	nextBatch.modelOutput, err = model.Forward(nextBatch.ctx, s.model, batch)
604
	if err != nil {
605
606
		err = fmt.Errorf("failed to build graph: %w", err)
		return
607
	}
608
609
	nextBatch.batchInputs = batchInputs
	nextBatch.batch = batch
610

611
612
613
614
615
616
617
618
619
620
621
622
	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
623
	logutil.Trace("computeBatch: waiting for inputs to be ready", "batchID", activeBatch.id)
624
	<-activeBatch.inputsReadyCh
Michael Yang's avatar
Michael Yang committed
625
	logutil.Trace("computeBatch: inputs are ready", "batchID", activeBatch.id)
626

627
628
629
	// 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
630
		logutil.Trace("computeBatch: outputs are ready", "batchID", activeBatch.id)
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
		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
647
	for i, seq := range s.seqs {
648
		iBatches[i] = -1
649
650
651
		if seq == nil {
			continue
		}
652
653
654
655
		// Skip over any newly added or skipped sequences
		if activeBatch.seqs[i] == nil {
			continue
		}
656

657
658
659
		// 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
660
			logutil.Trace("computeBatch: sequence replaced, discarding its results", "batchID", activeBatch.id, "seqIdx", i)
661
662
663
664
665
666
667
668
669
670
671
			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).
672
673
		if len(seq.pendingInputs) > 0 {
			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
674
			seq.pendingInputs = []*input.Input{}
675
676
		}

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

Jesse Gross's avatar
Jesse Gross committed
685
		seq.numPredicted++
686
687
688
689
690
691
692
693
694
		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
695
	activeBatch.batch.Inputs.FromInts(batchInputs)
696
697
	activeBatch.ctx.ComputeWithNotify(
		func() {
Michael Yang's avatar
Michael Yang committed
698
			logutil.Trace("computeBatch: signaling computeStartedCh", "batchID", activeBatch.id)
699
700
701
			activeBatch.computeStartedCh <- struct{}{}
		},
		activeBatch.modelOutput)
Michael Yang's avatar
Michael Yang committed
702
703

	outputs := activeBatch.modelOutput.Floats()
704
	t := time.Now()
705

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

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

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

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

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

		// sample a token
731
732
		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
733
		token, err := seq.sampler.Sample(outputs[iBatches[i]*vocabSize : (iBatches[i]+1)*vocabSize])
Jesse Gross's avatar
Jesse Gross committed
734
		if err != nil {
Michael Yang's avatar
Michael Yang committed
735
			panic("failed to sample token")
Jesse Gross's avatar
Jesse Gross committed
736
		}
737

738
739
		nextBatchTokens[i].Token = token

740
		// if it's an end of sequence token, break
Jesse Gross's avatar
Jesse Gross committed
741
		if s.model.(model.TextProcessor).Is(token, model.SpecialEOS) {
742
743
744
			// 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
745
			logutil.Trace("computeBatch: EOS", "batchID", activeBatch.id, "seqIdx", i)
746
			s.removeSequence(i, llm.DoneReasonStop)
747
748
749
			continue
		}

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

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

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

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

779
			seq.cache.Inputs = seq.cache.Inputs[:tokenLen]
780

781
			s.removeSequence(i, llm.DoneReasonStop)
782
783
784
			continue
		}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

				return
			}
		}
	}
}

Michael Yang's avatar
Michael Yang committed
936
func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
Michael Yang's avatar
Michael Yang committed
937
	if pooling.Type(s.model.Backend().Config().Uint("pooling_type")) == pooling.TypeNone {
Michael Yang's avatar
Michael Yang committed
938
939
940
941
942
943
944
945
946
947
948
		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")
949
950
	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{
		embedding: true,
951
952
953
954
955

		// 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,
956
	})
Michael Yang's avatar
Michael Yang committed
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
996
997
	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{
998
		Embedding: <-seq.embedding,
Michael Yang's avatar
Michael Yang committed
999
1000
1001
1002
1003
	}); err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
	}
}

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

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

1018
	var err error
1019
1020
1021
1022
	inputs := make([]*input.Input, s.batchSize)
	for i := range inputs {
		inputs[i] = &input.Input{}
	}
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
1062
1063
	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 {
1064
				newInputs := make([]*input.Input, s.batchSize)
1065
				copy(newInputs, inputs)
1066
1067
1068
				for i := len(inputs); i < s.batchSize; i++ {
					newInputs[i] = &input.Input{}
				}
1069
1070
1071
1072
1073
				inputs = newInputs
			}
		}
	}

1074
1075
	var batch input.Batch

1076
	batchInputs := make([]int32, len(inputs))
1077
1078
	batch.Positions = make([]int32, len(inputs))
	batch.Sequences = make([]int, len(inputs))
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
	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})
		}

1089
1090
1091
		batch.Positions[i] = int32(i)
	}

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

	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
	}

1108
	ctx.Forward(t).Reserve()
1109
1110

	return nil
1111
}
1112

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

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

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

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

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

1165
1166
1167
	return s.reserveWorstCaseGraph()
}

Jesse Gross's avatar
Jesse Gross committed
1168
1169
1170
1171
1172
1173
1174
// 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
1175
	}
Jesse Gross's avatar
Jesse Gross committed
1176
}
1177

Jesse Gross's avatar
Jesse Gross committed
1178
1179
1180
1181
// 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(),
1182
1183
1184
1185
		func(progress float32) {
			s.progress = progress
		})
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
1186
		panic(fmt.Errorf("failed to load model: %v", err))
1187
1188
	}

1189
	s.status = llm.ServerStatusReady
1190
1191
1192
	s.ready.Done()
}

Jesse Gross's avatar
Jesse Gross committed
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
1277
1278
// 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
	}
}

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

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

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

1341
1342
1343
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

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

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

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

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

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

1380
	return nil
1381
}