sched.go 31 KB
Newer Older
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1
2
3
4
5
6
7
package server

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
Michael Yang's avatar
int  
Michael Yang committed
8
	"os"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
9
	"reflect"
10
	"runtime"
11
	"slices"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
	"sort"
Michael Yang's avatar
int  
Michael Yang committed
13
	"strconv"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
14
15
16
17
18
	"strings"
	"sync"
	"time"

	"github.com/ollama/ollama/api"
19
	"github.com/ollama/ollama/discover"
20
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
21
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
22
	"github.com/ollama/ollama/fs/ggml"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
23
	"github.com/ollama/ollama/llm"
24
	"github.com/ollama/ollama/types/model"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
25
26
27
28
29
30
)

type LlmRequest struct {
	ctx             context.Context //nolint:containedctx
	model           *Model
	opts            api.Options
Daniel Hiltgen's avatar
Daniel Hiltgen committed
31
	origNumCtx      int // Track the initial ctx request
32
	sessionDuration *api.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
33
34
	successCh       chan *runnerRef
	errCh           chan error
35
	schedAttempts   uint
Daniel Hiltgen's avatar
Daniel Hiltgen committed
36
37
38
39
40
41
}

type Scheduler struct {
	pendingReqCh  chan *LlmRequest
	finishedReqCh chan *LlmRequest
	expiredCh     chan *runnerRef
42
	unloadedCh    chan any
Daniel Hiltgen's avatar
Daniel Hiltgen committed
43
44
45
46

	loaded   map[string]*runnerRef
	loadedMu sync.Mutex

Michael Yang's avatar
Michael Yang committed
47
48
	loadFn       func(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel int)
	newServerFn  func(gpus discover.GpuInfoList, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int) (llm.LlamaServer, error)
49
50
	getGpuFn     func() discover.GpuInfoList
	getCpuFn     func() discover.GpuInfoList
51
	reschedDelay time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
52
53
}

54
55
56
57
58
59
// Default automatic value for number of models we allow per GPU
// Model will still need to fit in VRAM, but loading many small models
// on a large GPU can cause stalling
var defaultModelsPerGPU = 3

// Default automatic value for parallel setting
60
// Model will still need to fit in VRAM.  If this setting won't fit
61
// we'll back off down to 1 to try to get it to fit
Devon Rifkin's avatar
Devon Rifkin committed
62
var defaultParallel = 2
63

Michael Yang's avatar
lint  
Michael Yang committed
64
var ErrMaxQueue = errors.New("server busy, please try again.  maximum pending requests exceeded")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
65
66

func InitScheduler(ctx context.Context) *Scheduler {
Michael Yang's avatar
int  
Michael Yang committed
67
	maxQueue := envconfig.MaxQueue()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
68
	sched := &Scheduler{
Michael Yang's avatar
int  
Michael Yang committed
69
70
71
		pendingReqCh:  make(chan *LlmRequest, maxQueue),
		finishedReqCh: make(chan *LlmRequest, maxQueue),
		expiredCh:     make(chan *runnerRef, maxQueue),
72
		unloadedCh:    make(chan any, maxQueue),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
73
74
		loaded:        make(map[string]*runnerRef),
		newServerFn:   llm.NewLlamaServer,
75
76
		getGpuFn:      discover.GetGPUInfo,
		getCpuFn:      discover.GetCPUInfo,
77
		reschedDelay:  250 * time.Millisecond,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
78
79
80
81
82
83
	}
	sched.loadFn = sched.load
	return sched
}

// context must be canceled to decrement ref count and release the runner
84
func (s *Scheduler) GetRunner(c context.Context, model *Model, opts api.Options, sessionDuration *api.Duration) (chan *runnerRef, chan error) {
85
86
87
88
	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
89
90
91
92
93
94
95
96
	req := &LlmRequest{
		ctx:             c,
		model:           model,
		opts:            opts,
		sessionDuration: sessionDuration,
		successCh:       make(chan *runnerRef),
		errCh:           make(chan error, 1),
	}
97

Daniel Hiltgen's avatar
Daniel Hiltgen committed
98
99
100
	select {
	case s.pendingReqCh <- req:
	default:
101
		req.errCh <- ErrMaxQueue
Daniel Hiltgen's avatar
Daniel Hiltgen committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
	}
	return req.successCh, req.errCh
}

// Returns immediately, spawns go routines for the scheduler which will shutdown when ctx is done
func (s *Scheduler) Run(ctx context.Context) {
	slog.Debug("starting llm scheduler")
	go func() {
		s.processPending(ctx)
	}()

	go func() {
		s.processCompleted(ctx)
	}()
}

func (s *Scheduler) processPending(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			slog.Debug("shutting down scheduler pending loop")
			return
		case pending := <-s.pendingReqCh:
			// Block other requests until we get this pending request running
126
			pending.schedAttempts++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
127
128
			if pending.origNumCtx == 0 {
				pending.origNumCtx = pending.opts.NumCtx
Daniel Hiltgen's avatar
Daniel Hiltgen committed
129
			}
130
131
132
133
134

			if pending.ctx.Err() != nil {
				slog.Debug("pending request cancelled or timed out, skipping scheduling")
				continue
			}
Michael Yang's avatar
Michael Yang committed
135
			numParallel := int(envconfig.NumParallel())
136
137
138
			// `mllama` is a snowflake and uses an encoder cache which cannot be used with num_parallel > 1
			// ref: https://github.com/ollama/ollama/issues/4165
			if slices.Contains(pending.model.Config.ModelFamilies, "mllama") && numParallel != 1 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
				numParallel = 1
140
				slog.Warn("mllama does not currently support parallel requests")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
141
			}
142

Daniel Hiltgen's avatar
Daniel Hiltgen committed
143
144
145
146
147
148
149
150
			for {
				var runnerToExpire *runnerRef
				s.loadedMu.Lock()
				runner := s.loaded[pending.model.ModelPath]
				loadedCount := len(s.loaded)
				s.loadedMu.Unlock()
				if runner != nil {
					if runner.needsReload(ctx, pending) {
151
						slog.Debug("reloading", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
152
153
154
155
156
157
						runnerToExpire = runner
					} else {
						// Runner is usable, return it
						pending.useLoadedRunner(runner, s.finishedReqCh)
						break
					}
Michael Yang's avatar
Michael Yang committed
158
				} else if envconfig.MaxRunners() > 0 && loadedCount >= int(envconfig.MaxRunners()) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
159
					slog.Debug("max runners achieved, unloading one to make room", "runner_count", loadedCount)
160
					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
161
				} else {
162
					// Either no models are loaded or below envconfig.MaxRunners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
163
					// Get a refreshed GPU list
164
					var gpus discover.GpuInfoList
165
166
167
168
169
					if pending.opts.NumGPU == 0 {
						gpus = s.getCpuFn()
					} else {
						gpus = s.getGpuFn()
					}
170

Michael Yang's avatar
int  
Michael Yang committed
171
					if envconfig.MaxRunners() <= 0 {
172
173
174
175
176
177
178
179
180
181
182
						// No user specified MaxRunners, so figure out what automatic setting to use
						// If all GPUs have reliable free memory reporting, defaultModelsPerGPU * the number of GPUs
						// if any GPU has unreliable free memory reporting, 1x the number of GPUs
						allReliable := true
						for _, gpu := range gpus {
							if gpu.UnreliableFreeMemory {
								allReliable = false
								break
							}
						}
						if allReliable {
Michael Yang's avatar
int  
Michael Yang committed
183
184
							// HACK
							os.Setenv("OLLAMA_MAX_LOADED_MODELS", strconv.Itoa(defaultModelsPerGPU*len(gpus)))
185
							slog.Debug("updating default concurrency", "OLLAMA_MAX_LOADED_MODELS", envconfig.MaxRunners(), "gpu_count", len(gpus))
186
						} else {
Michael Yang's avatar
int  
Michael Yang committed
187
188
							// HACK
							os.Setenv("OLLAMA_MAX_LOADED_MODELS", strconv.Itoa(len(gpus)))
189
190
191
							slog.Info("one or more GPUs detected that are unable to accurately report free memory - disabling default concurrency")
						}
					}
192

193
					// Load model for fitting
194
					ggml, err := llm.LoadModel(pending.model.ModelPath, 0)
195
196
197
198
					if err != nil {
						pending.errCh <- err
						break
					}
199

200
					// Embedding models should always be loaded with parallel=1
201
					if pending.model.CheckCapabilities(model.CapabilityCompletion) != nil {
202
203
204
						numParallel = 1
					}

205
206
					// Evaluate if the model will fit in the available system memory, or if we should unload a model first
					if len(gpus) == 1 && gpus[0].Library == "cpu" {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
207
208
209
210
211
						// simplifying assumption of defaultParallel when in CPU mode
						if numParallel <= 0 {
							numParallel = defaultParallel
						}

212
213
						pending.opts.NumCtx = pending.origNumCtx * numParallel

214
215
						if loadedCount == 0 {
							slog.Debug("cpu mode with first model, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
216
							s.loadFn(pending, ggml, gpus, numParallel)
217
218
219
220
221
							break
						}
						runnerToExpire = s.maybeFindCPURunnerToUnload(pending, ggml, gpus)
						if runnerToExpire == nil {
							slog.Debug("cpu mode with available system memory or first model, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
222
							s.loadFn(pending, ggml, gpus, numParallel)
223
224
225
226
227
							break
						}
						// else we need to expire a runner
					} else if loadedCount == 0 {
						// No models loaded. Load the model but prefer the best fit.
228
						slog.Debug("loading first model", "model", pending.model.ModelPath)
229
						g := pickBestFullFitByLibrary(pending, ggml, gpus, &numParallel)
230
231
						if g != nil {
							gpus = g
232
233
234
						} else {
							// Only allow partial loads when this is the first model
							gpus = pickBestPartialFitByLibrary(pending, ggml, gpus, &numParallel)
235
						}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
236
						s.loadFn(pending, ggml, gpus, numParallel)
237
238
239
						break
					}

240
					if runnerToExpire == nil {
241
242
243
244
245
246
						// More than one loaded model, so we have to see if the
						// new one fits
						//
						// We want to avoid loading on any GPUs that have other
						// models still loading on them to avoid potential races
						// with VRAM consumption ramping up during load
Daniel Hiltgen's avatar
Daniel Hiltgen committed
247
						availGpus := s.filterGPUsWithoutLoadingModels(gpus)
248

249
						// Update free memory from currently loaded models
250
						s.updateFreeSpace(availGpus)
251
						fitGpus := pickBestFullFitByLibrary(pending, ggml, availGpus, &numParallel)
252
						if fitGpus != nil {
253
							slog.Debug("new model fits with existing models, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
254
							s.loadFn(pending, ggml, fitGpus, numParallel)
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
							break
						}

						// We couldn't find a set of GPUs to fully load the new
						// model. If no other models are loading (both GPU lists
						// are the same) then we need to unload another model to
						// make room
						if len(availGpus) < len(gpus) {
							// There are other requests pending, and this one
							// needs more time, so put it on the back of the
							// queue so that we might satisfy other pending
							// requests that aren't blocked
							go func() {
								// Process in a go routine to avoid deadlocking
								// the scheduler if our queue is full
								slog.Debug("delaying scheduling while other models finish loading", "attempts", pending.schedAttempts, "model", pending.model.ModelPath)
								time.Sleep(s.reschedDelay)
								s.pendingReqCh <- pending
							}()
274
275
276
							break
						}
						runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
277
278
279
280
281
282
283
284
285
286
					}
				}

				if runnerToExpire == nil {
					// Shouildn't happen
					slog.Error("runner to expire was nil!")
					continue
				}
				// Trigger an expiration to unload once it's done
				runnerToExpire.refMu.Lock()
287
				slog.Debug("resetting model to expire immediately to make room", "runner", runnerToExpire, "refCount", runnerToExpire.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
288
289
290
291
292
293
294
295
296
297
298
299
				if runnerToExpire.expireTimer != nil {
					runnerToExpire.expireTimer.Stop()
					runnerToExpire.expireTimer = nil
				}
				runnerToExpire.sessionDuration = 0
				if runnerToExpire.refCount <= 0 {
					s.expiredCh <- runnerToExpire
				}
				runnerToExpire.refMu.Unlock()
				// Wait for the unload to happen
				// Note: at this point we're queueing up all incoming requests, even if they were for
				// a different model that's loaded and not scheduled to be removed.
300
				slog.Debug("waiting for pending requests to complete and unload to occur", "runner", runnerToExpire)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
301
302
303
304
305
				select {
				case <-ctx.Done():
					slog.Debug("shutting down scheduler pending loop")
					return
				case <-s.unloadedCh:
306
					slog.Debug("unload completed", "runner", runnerToExpire)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
					continue
				}
			}
		case <-s.unloadedCh:
			// An unload request when there are no pending request can be ignored
			slog.Debug("ignoring unload event with no pending requests")
		}
	}
}

func (s *Scheduler) processCompleted(ctx context.Context) {
	// Process completed requests, expired timers, and unloading models
	for {
		select {
		case <-ctx.Done():
			slog.Debug("shutting down scheduler completed loop")
			return
		case finished := <-s.finishedReqCh:
			s.loadedMu.Lock()
			runner := s.loaded[finished.model.ModelPath]
			s.loadedMu.Unlock()
			if runner == nil {
329
				slog.Error("finished request signal received after model unloaded", "modelPath", finished.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
330
331
332
333
334
335
				continue
			}
			runner.refMu.Lock()
			runner.refCount--
			if runner.refCount <= 0 {
				if runner.sessionDuration <= 0 {
336
					slog.Debug("runner with zero duration has gone idle, expiring to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
337
338
339
340
341
342
					if runner.expireTimer != nil {
						runner.expireTimer.Stop()
						runner.expireTimer = nil
					}
					s.expiredCh <- runner
				} else if runner.expireTimer == nil {
343
					slog.Debug("runner with non-zero duration has gone idle, adding timer", "runner", runner, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
344
					runner.expireTimer = time.AfterFunc(runner.sessionDuration, func() {
345
						slog.Debug("timer expired, expiring to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
346
347
348
349
						runner.refMu.Lock()
						defer runner.refMu.Unlock()
						if runner.expireTimer != nil {
							runner.expireTimer.Stop()
350
							runner.expireTimer = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
351
352
353
						}
						s.expiredCh <- runner
					})
354
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
355
				} else {
356
					slog.Debug("runner with non-zero duration has gone idle, resetting timer", "runner", runner, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
357
					runner.expireTimer.Reset(runner.sessionDuration)
358
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
359
360
				}
			}
361
			slog.Debug("after processing request finished event", "runner", runner, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
362
363
			runner.refMu.Unlock()
		case runner := <-s.expiredCh:
364
			slog.Debug("runner expired event received", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
365
366
			runner.refMu.Lock()
			if runner.refCount > 0 {
367
				slog.Debug("expired event with positive ref count, retrying", "runner", runner, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
368
369
370
371
372
373
374
375
376
377
				go func(runner *runnerRef) {
					// We can't unload yet, but want to as soon as the current request completes
					// So queue up another expired event
					time.Sleep(10 * time.Millisecond)
					s.expiredCh <- runner
				}(runner)
				runner.refMu.Unlock()
				continue
			}

378
			s.loadedMu.Lock()
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
			slog.Debug("got lock to unload expired event", "runner", runner)
			runnerToUnload := s.loaded[runner.modelPath]
			if runnerToUnload == nil {
				// If runnerToUnload is nil, we already processed an event and
				// unloaded it. This double unload can happen if the initial
				// request is canceled and we're trying to load another model
				// that requires this one to be evicted, or the settings change
				// and require a reload
				s.loadedMu.Unlock()
				runner.refMu.Unlock()
				slog.Debug("duplicate expired event, ignoring", "runner", runner)
			} else {
				slog.Debug("starting background wait for VRAM recovery", "runner", runner)
				finished := runner.waitForVRAMRecovery()
				runner.unload()
				delete(s.loaded, runner.modelPath)
				s.loadedMu.Unlock()
				slog.Debug("runner terminated and removed from list, blocking for VRAM recovery", "runner", runner)
				<-finished
				runner.refMu.Unlock()
				slog.Debug("sending an unloaded event", "runner", runner)
				s.unloadedCh <- struct{}{}
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
402
403
404
405
406
407
408
409
410
411
412
		}
	}
}

// Complete the pending request and send the runner back to the requester
// Wires up a finished event after the request context is completed
// Updates session duration, and resets expiration timer
func (pending *LlmRequest) useLoadedRunner(runner *runnerRef, finished chan *LlmRequest) {
	runner.refMu.Lock()
	defer runner.refMu.Unlock()
	runner.refCount++
413
414
415
416
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
417
418
419
	if pending.sessionDuration != nil {
		runner.sessionDuration = pending.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
420
421
422
	pending.successCh <- runner
	go func() {
		<-pending.ctx.Done()
423
		slog.Debug("context for request finished", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
424
425
426
427
		finished <- pending
	}()
}

Michael Yang's avatar
Michael Yang committed
428
func (s *Scheduler) load(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel int) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
429
430
431
	if numParallel < 1 {
		numParallel = 1
	}
Michael Yang's avatar
Michael Yang committed
432
	sessionDuration := envconfig.KeepAlive()
433
434
435
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Michael Yang's avatar
Michael Yang committed
436
	llama, err := s.newServerFn(gpus, req.model.ModelPath, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, numParallel)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
437
438
439
440
	if err != nil {
		// some older models are not compatible with newer versions of llama.cpp
		// show a generalized compatibility error until there is a better way to
		// check for model compatibility
Michael Yang's avatar
Michael Yang committed
441
		if errors.Is(err, ggml.ErrUnsupportedFormat) || strings.Contains(err.Error(), "failed to load model") {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
442
443
444
445
446
447
			err = fmt.Errorf("%v: this model may be incompatible with your version of Ollama. If you previously pulled this model, try updating it by running `ollama pull %s`", err, req.model.ShortName)
		}
		slog.Info("NewLlamaServer failed", "model", req.model.ModelPath, "error", err)
		req.errCh <- err
		return
	}
448
449
450
451
452
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
453
		sessionDuration: sessionDuration,
454
455
456
457
		gpus:            gpus,
		estimatedVRAM:   llama.EstimatedVRAM(),
		estimatedTotal:  llama.EstimatedTotal(),
		loading:         true,
458
		pid:             llama.Pid(),
459
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
460
	runner.numParallel = numParallel
461
	runner.refMu.Lock() // hold lock until running or aborted
462

Daniel Hiltgen's avatar
Daniel Hiltgen committed
463
	s.loadedMu.Lock()
464
465
466
467
468
469
470
	if oldRunner, ok := s.loaded[req.model.ModelPath]; ok {
		// Shouldn't happen, but safeguard against leaking a runner
		slog.Warn("model was still loaded", "old_runner", oldRunner, "new_runner", runner)
		oldRunner.refMu.Lock()
		oldRunner.unload()
		oldRunner.refMu.Unlock()
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
471
472
473
474
475
476
477
478
479
	s.loaded[req.model.ModelPath] = runner
	slog.Info("loaded runners", "count", len(s.loaded))
	s.loadedMu.Unlock()

	go func() {
		defer runner.refMu.Unlock()
		if err = llama.WaitUntilRunning(req.ctx); err != nil {
			slog.Error("error loading llama server", "error", err)
			req.errCh <- err
480
			slog.Debug("triggering expiration for failed load", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
481
482
483
			s.expiredCh <- runner
			return
		}
484
		slog.Debug("finished setting up", "runner", runner)
485
486
487
		if runner.pid < 0 {
			runner.pid = llama.Pid()
		}
488
		runner.refCount++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
489
490
491
492
493
494
495
496
497
498
		runner.loading = false
		go func() {
			<-req.ctx.Done()
			slog.Debug("context for request finished")
			s.finishedReqCh <- req
		}()
		req.successCh <- runner
	}()
}

499
func (s *Scheduler) updateFreeSpace(allGpus discover.GpuInfoList) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
500
501
502
503
504
505
	type predKey struct {
		Library string
		ID      string
	}
	predMap := map[predKey]uint64{} // Sum up the total predicted usage per GPU for all runners
	s.loadedMu.Lock()
506
	runners := make([]*runnerRef, 0, len(s.loaded))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
507
	for _, r := range s.loaded {
508
509
510
511
		runners = append(runners, r)
	}
	s.loadedMu.Unlock()
	for _, r := range runners {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
512
513
514
		r.refMu.Lock()
		if r.llama != nil {
			for _, gpu := range allGpus {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
515
				predMap[predKey{gpu.Library, gpu.ID}] += r.llama.EstimatedVRAMByGPU(gpu.ID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
			}
		} else {
			slog.Warn("unexpected nil runner reference, memory prediction may be incorrect")
		}
		r.refMu.Unlock()
	}

	// Now that we've summed up all the GPU usage predictions across all the loaded runners, update the gpu list
	for i := range allGpus {
		if p, ok := predMap[predKey{allGpus[i].Library, allGpus[i].ID}]; ok {
			slog.Debug("gpu reported", "gpu", allGpus[i].ID, "library", allGpus[i].Library, "available", format.HumanBytes2(allGpus[i].FreeMemory))
			if p > allGpus[i].TotalMemory {
				// Shouldn't happen
				slog.Warn("predicted usage exceeds VRAM", "gpu", allGpus[i].ID, "totalMemory", allGpus[i].TotalMemory, "predicted", p)
				allGpus[i].FreeMemory = 0
			} else if (allGpus[i].TotalMemory - p) < allGpus[i].FreeMemory { // predicted free is smaller than reported free, use it
				// TODO maybe we should just always trust our numbers, since cuda's free memory reporting is laggy
				// and we might unload models we didn't actually need to.  The risk is if some other GPU intensive app is loaded
534
				// after we start our first runner, then we'll never account for that, so picking the smallest free value seems prudent.
Daniel Hiltgen's avatar
Daniel Hiltgen committed
535
536
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
537
538
539
540
541
542
543
544
545
			slog.Info("updated VRAM based on existing loaded models", "gpu", allGpus[i].ID, "library", allGpus[i].Library, "total", format.HumanBytes2(allGpus[i].TotalMemory), "available", format.HumanBytes2(allGpus[i].FreeMemory))
		}
	}
}

// While models are loading the VRAM consumption numbers will be indeterminate, so we have
// to avoid scheduling another model on the same GPU(s) that haven't stabilized.
// This routine returns the set of GPUs that do not have an active loading model.
// If all GPUs have loading models, an empty list will be returned (not a single CPU entry)
546
547
func (s *Scheduler) filterGPUsWithoutLoadingModels(allGpus discover.GpuInfoList) discover.GpuInfoList {
	ret := append(discover.GpuInfoList{}, allGpus...)
548
549
550
551
552
553
554
555
556
557
558
559
560
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
	for _, runner := range s.loaded {
		if runner.loading {
			slog.Debug("overlapping loads detected", "gpus", runner.gpus, "model", runner.modelPath)
			for _, busyGPU := range runner.gpus {
				for i := range ret {
					if ret[i].ID == busyGPU.ID {
						ret = append(ret[:i], ret[i+1:]...)
						break
					}
				}
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
561
562
		}
	}
563
	return ret
Daniel Hiltgen's avatar
Daniel Hiltgen committed
564
565
}

566
// TODO consolidate sched_types.go
Daniel Hiltgen's avatar
Daniel Hiltgen committed
567
type runnerRef struct {
568
	refMu    sync.Mutex
Daniel Hiltgen's avatar
Daniel Hiltgen committed
569
570
	refCount uint // prevent unloading if > 0

571
	llama          llm.LlamaServer
572
	pid            int
573
574
	loading        bool                 // True only during initial load, then false forever
	gpus           discover.GpuInfoList // Recorded at time of provisioning
575
576
	estimatedVRAM  uint64
	estimatedTotal uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
577
578
579

	sessionDuration time.Duration
	expireTimer     *time.Timer
580
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
581

Daniel Hiltgen's avatar
Daniel Hiltgen committed
582
583
584
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
585
586
587
588
589
	*api.Options
}

// The refMu must already be held when calling unload
func (runner *runnerRef) unload() {
590
591
592
593
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
594
595
596
	if runner.llama != nil {
		runner.llama.Close()
	}
597
	runner.model = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
598
599
600
601
602
603
604
605
606
	runner.llama = nil
	runner.Options = nil
	runner.gpus = nil
}

func (runner *runnerRef) needsReload(ctx context.Context, req *LlmRequest) bool {
	slog.Debug("evaluating already loaded", "model", req.model.ModelPath)
	runner.refMu.Lock()
	defer runner.refMu.Unlock()
607

Daniel Hiltgen's avatar
Daniel Hiltgen committed
608
609
610
611
	timeout := 10 * time.Second
	if runner.loading {
		timeout = 2 * time.Minute // Initial load can take a long time for big models on slow systems...
	}
612

613
614
615
616
	if runner.Options == nil {
		return true
	}

617
618
619
620
621
622
623
624
	// Don't reload runner if num_gpu=-1 was provided
	optsExisting := runner.Options.Runner
	optsNew := req.opts.Runner
	if optsNew.NumGPU < 0 {
		optsExisting.NumGPU = -1
		optsNew.NumGPU = -1
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
625
626
627
	// Normalize the NumCtx for parallelism
	optsExisting.NumCtx = optsExisting.NumCtx / runner.numParallel

628
	ctx, cancel := context.WithTimeout(ctx, timeout)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
629
	defer cancel()
630
631
	if !reflect.DeepEqual(runner.model.AdapterPaths, req.model.AdapterPaths) || // have the adapters changed?
		!reflect.DeepEqual(runner.model.ProjectorPaths, req.model.ProjectorPaths) || // have the projectors changed?
Daniel Hiltgen's avatar
Daniel Hiltgen committed
632
633
634
635
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
636

Daniel Hiltgen's avatar
Daniel Hiltgen committed
637
638
639
	return false
}

640
641
642
643
644
645
646
647
// Free memory reporting on GPUs can lag for a while even after the runner
// exits, so we have to keep checking until we see the available memory recover,
// otherwise subsequent model loads will get far less layers loaded or worse
// case, may completely fall back to CPU mode.
// This routine must be called before the runner unloads so it can establish
// a before and after GPU memory allocation.  The returned channel
// will be notified when we're done waiting, or have timed out and should
// proceed anyway
648
649
func (runner *runnerRef) waitForVRAMRecovery() chan any {
	finished := make(chan any, 1)
650

651
652
	// CPU or Metal don't need checking, so no waiting required
	// windows can page VRAM, only cuda currently can report accurate used vram usage
653
654
	if len(runner.gpus) == 0 ||
		(len(runner.gpus) == 1 && (runner.gpus[0].Library == "cpu" || runner.gpus[0].Library == "metal")) ||
655
		(runtime.GOOS == "windows" && runner.gpus[0].Library != "cuda") {
656
		finished <- struct{}{}
657
		slog.Debug("no need to wait for VRAM recovery", "runner", runner)
658
659
660
661
662
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
663
	gpusBefore := discover.GetGPUInfo()
664
665
666
667
668
669
670
671
672
673
674
675
	var totalMemoryBefore, freeMemoryBefore uint64
	for _, gpu := range gpusBefore {
		totalMemoryBefore += gpu.TotalMemory
		freeMemoryBefore += gpu.FreeMemory
	}
	go func() {
		expiresAt := start.Add(5 * time.Second) // typical convergence is 0.5-1.5s
		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
			<-ticker.C
			if time.Now().After(expiresAt) {
676
				slog.Warn("gpu VRAM usage didn't recover within timeout", "seconds", time.Since(start).Seconds(), "runner", runner)
677
678
679
680
				finished <- struct{}{}
			}

			// Query GPUs, look for free to go back up
681
			gpusNow := discover.GetGPUInfo()
682
683
684
685
686
687
688
			var totalMemoryNow, freeMemoryNow uint64
			for _, gpu := range gpusNow {
				totalMemoryNow += gpu.TotalMemory
				freeMemoryNow += gpu.FreeMemory
			}
			// If we're within ~80% of the estimated memory usage recovered, bail out
			if float32(freeMemoryNow-freeMemoryBefore) > float32(runner.estimatedVRAM)*0.8 {
689
				slog.Debug(fmt.Sprintf("gpu VRAM free memory converged after %0.2f seconds", time.Since(start).Seconds()), "runner", runner)
690
691
692
693
694
695
696
697
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
func (runner *runnerRef) LogValue() slog.Value {
	if runner == nil {
		return slog.StringValue("nil")
	}
	attrs := []slog.Attr{}
	if runner.model != nil {
		attrs = append(attrs, slog.String("name", runner.model.Name))
	}
	if len(runner.gpus) > 0 {
		attrs = append(attrs,
			slog.String("inference", runner.gpus[0].Library),
			slog.Int("devices", len(runner.gpus)),
		)
	}
	attrs = append(attrs,
		slog.String("size", format.HumanBytes2(runner.estimatedTotal)),
		slog.String("vram", format.HumanBytes2(runner.estimatedVRAM)),
		slog.Int("parallel", runner.numParallel),
		slog.Int("pid", runner.pid),
		slog.String("model", runner.modelPath),
	)
719
720
721
	if runner.Options != nil {
		attrs = append(attrs, slog.Int("num_ctx", runner.Options.NumCtx))
	}
722
723
724
	return slog.GroupValue(attrs...)
}

725
726
727
728
729
730
731
732
733
734
735
736
737
type ByDurationAndName []*runnerRef

func (a ByDurationAndName) Len() int      { return len(a) }
func (a ByDurationAndName) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByDurationAndName) Less(i, j int) bool {
	// Primary sort by session duration (uint64 to handle negatives)
	d1 := uint64(a[i].sessionDuration)
	d2 := uint64(a[j].sessionDuration)
	if d1 != d2 {
		return d1 < d2
	}
	// Secondary sort by model path lex order
	return a[i].modelPath < a[j].modelPath
Daniel Hiltgen's avatar
Daniel Hiltgen committed
738
739
740
741
742
743
744
745
}

// TODO - future consideration to pick runners based on size
// type BySize []*runnerRef
// func (a BySize) Len() int           { return len(a) }
// func (a BySize) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
// func (a BySize) Less(i, j int) bool { return a[i].estimatedVRAM < a[j].estimatedVRAM }

746
747
// pickBestFullFitByLibrary will try to find the optimal placement of the model in the available GPUs where the model fully fits
// The list of GPUs returned will always be the same brand (library)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
748
// If the model can not be fit fully within the available GPU(s) nil is returned
749
// If numParallel is <= 0, this will attempt try to optimize parallelism based on available VRAM, and adjust
Daniel Hiltgen's avatar
Daniel Hiltgen committed
750
// opts.NumCtx accordingly
Michael Yang's avatar
Michael Yang committed
751
func pickBestFullFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel *int) discover.GpuInfoList {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
752
	var estimatedVRAM uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
753
754
755
756

	var numParallelToTry []int
	if *numParallel <= 0 {
		// If no specific parallel setting was provided, try larger then smaller, always end with 1
757
		numParallelToTry = append(numParallelToTry, defaultParallel, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
758
759
760
761
	} else {
		numParallelToTry = []int{*numParallel}
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
762
763
	for _, gl := range gpus.ByLibrary() {
		var ok bool
764
		sgl := append(make(discover.GpuInfoList, 0, len(gl)), gl...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
765
766

		// TODO - potentially sort by performance capability, existing models loaded, etc.
Daniel Hiltgen's avatar
Daniel Hiltgen committed
767
		// TODO - Eliminate any GPUs that already have envconfig.MaxRunners loaded on them
Daniel Hiltgen's avatar
Daniel Hiltgen committed
768
		// Note: at present, this will favor more VRAM over faster GPU speed in mixed setups
769
		sort.Sort(sort.Reverse(discover.ByFreeMemory(sgl)))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
770
771

		// First attempt to fit the model into a single GPU
Daniel Hiltgen's avatar
Daniel Hiltgen committed
772
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
773
			req.opts.NumCtx = req.origNumCtx * p
Michael Yang's avatar
bool  
Michael Yang committed
774
			if !envconfig.SchedSpread() {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
775
				for _, g := range sgl {
776
					if ok, estimatedVRAM = llm.PredictServerFit([]discover.GpuInfo{g}, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, p); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
777
778
						slog.Info("new model will fit in available VRAM in single GPU, loading", "model", req.model.ModelPath, "gpu", g.ID, "parallel", p, "available", g.FreeMemory, "required", format.HumanBytes2(estimatedVRAM))
						*numParallel = p
779
						return []discover.GpuInfo{g}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
780
					}
781
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
782
783
784
785
786
787
788
789
			}
		}

		// TODO future refinements
		// - if multiple Libraries, see if any single GPU in any Library will fit
		// - try subsets of GPUs instead of just falling back to 1 or all in a family

		// Now try all the GPUs
Daniel Hiltgen's avatar
Daniel Hiltgen committed
790
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
791
			req.opts.NumCtx = req.origNumCtx * p
792
			if ok, estimatedVRAM = llm.PredictServerFit(sgl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, p); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
793
794
795
796
				slog.Info("new model will fit in available VRAM, loading", "model", req.model.ModelPath, "library", sgl[0].Library, "parallel", p, "required", format.HumanBytes2(estimatedVRAM))
				*numParallel = p
				return sgl
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
797
798
799
800
801
		}
	}
	return nil
}

802
// If multiple Libraries are detected, pick the Library which loads the most layers for the model
Michael Yang's avatar
Michael Yang committed
803
func pickBestPartialFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel *int) discover.GpuInfoList {
804
805
	if *numParallel <= 0 {
		*numParallel = 1
Richard Lyons's avatar
Richard Lyons committed
806
807
		req.opts.NumCtx = req.origNumCtx
	}
808
809
810
811
812
813
814
	byLibrary := gpus.ByLibrary()
	if len(byLibrary) <= 1 {
		return gpus
	}
	var bestEstimate uint64
	var bestFit int
	for i, gl := range byLibrary {
815
		_, estimatedVRAM := llm.PredictServerFit(gl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, *numParallel)
816
817
818
819
820
821
822
823
		if estimatedVRAM > bestEstimate {
			bestEstimate = estimatedVRAM
			bestFit = i
		}
	}
	return byLibrary[bestFit]
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
824
// findRunnerToUnload finds a runner to unload to make room for a new model
825
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
826
827
828
829
830
831
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
832
833
834
835
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
836
837
838

	// In the future we can enhance the algorithm to be smarter about picking the optimal runner to unload
	// e.g., if we have multiple options, will one make room for the request?
839
	sort.Sort(ByDurationAndName(runnerList))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
840
841
842
843
844
845
846

	// First try to find a runner that's already idle
	for _, runner := range runnerList {
		runner.refMu.Lock()
		rc := runner.refCount
		runner.refMu.Unlock()
		if rc == 0 {
847
			slog.Debug("found an idle runner to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
848
849
850
851
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
852
	slog.Debug("no idle runners, picking the shortest duration", "runner_count", len(runnerList), "runner", runnerList[0])
Daniel Hiltgen's avatar
Daniel Hiltgen committed
853
854
855
856
857
858
859
860
861
862
863
864
865
	return runnerList[0]
}

func (s *Scheduler) unloadAllRunners() {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
	for model, runner := range s.loaded {
		if runner.llama != nil {
			slog.Debug("shutting down runner", "model", model)
			runner.llama.Close()
		}
	}
}
866

Patrick Devine's avatar
Patrick Devine committed
867
868
869
func (s *Scheduler) expireRunner(model *Model) {
	s.loadedMu.Lock()
	runner, ok := s.loaded[model.ModelPath]
870
	s.loadedMu.Unlock()
Patrick Devine's avatar
Patrick Devine committed
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
	if ok {
		runner.refMu.Lock()
		runner.expiresAt = time.Now()
		if runner.expireTimer != nil {
			runner.expireTimer.Stop()
			runner.expireTimer = nil
		}
		runner.sessionDuration = 0
		if runner.refCount <= 0 {
			s.expiredCh <- runner
		}
		runner.refMu.Unlock()
	}
}

886
887
// If other runners are loaded, make sure the pending request will fit in system memory
// If not, pick a runner to unload, else return nil and the request can be loaded
Michael Yang's avatar
Michael Yang committed
888
func (s *Scheduler) maybeFindCPURunnerToUnload(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList) *runnerRef {
889
	slog.Debug("evaluating if CPU model load will fit in available system memory")
890
	estimate := llm.EstimateGPULayers(gpus, f, req.model.ProjectorPaths, req.opts, req.opts.NumCtx/req.origNumCtx)
891
892
893
894
895
896
897
898
899
	if estimate.TotalSize <= gpus[0].FreeMemory {
		slog.Debug("cpu inference mode, model fits in available system memory", "model", format.HumanBytes2(estimate.TotalSize), "available", format.HumanBytes2(gpus[0].FreeMemory))
		return nil
	}

	// TODO - optimization: try to find CPU only runners first, or partial offloads with enough in system memory to make room

	return s.findRunnerToUnload()
}