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

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"reflect"
9
	"runtime"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
10
11
12
13
14
15
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/ollama/ollama/api"
16
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
17
18
19
20
21
22
23
24
25
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/gpu"
	"github.com/ollama/ollama/llm"
)

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

type Scheduler struct {
	pendingReqCh  chan *LlmRequest
	finishedReqCh chan *LlmRequest
	expiredCh     chan *runnerRef
	unloadedCh    chan interface{}

	loaded   map[string]*runnerRef
	loadedMu sync.Mutex

Daniel Hiltgen's avatar
Daniel Hiltgen committed
42
43
	loadFn       func(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList, numParallel int)
	newServerFn  func(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options, numParallel int) (llm.LlamaServer, error)
44
45
46
	getGpuFn     func() gpu.GpuInfoList
	getCpuFn     func() gpu.GpuInfoList
	reschedDelay time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
47
48
}

49
50
51
52
53
54
55
56
57
58
// 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
// Model will still need to fit in VRAM.  If this setting wont fit
// we'll back off down to 1 to try to get it to fit
var defaultParallel = 4

59
var ErrMaxQueue = fmt.Errorf("server busy, please try again.  maximum pending requests exceeded")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
60
61
62

func InitScheduler(ctx context.Context) *Scheduler {
	sched := &Scheduler{
63
64
65
66
		pendingReqCh:  make(chan *LlmRequest, envconfig.MaxQueuedRequests),
		finishedReqCh: make(chan *LlmRequest, envconfig.MaxQueuedRequests),
		expiredCh:     make(chan *runnerRef, envconfig.MaxQueuedRequests),
		unloadedCh:    make(chan interface{}, envconfig.MaxQueuedRequests),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
67
68
69
		loaded:        make(map[string]*runnerRef),
		newServerFn:   llm.NewLlamaServer,
		getGpuFn:      gpu.GetGPUInfo,
70
		getCpuFn:      gpu.GetCPUInfo,
71
		reschedDelay:  250 * time.Millisecond,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
72
73
74
75
76
77
	}
	sched.loadFn = sched.load
	return sched
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
83
84
85
86
87
88
89
90
	req := &LlmRequest{
		ctx:             c,
		model:           model,
		opts:            opts,
		sessionDuration: sessionDuration,
		successCh:       make(chan *runnerRef),
		errCh:           make(chan error, 1),
	}
91

Daniel Hiltgen's avatar
Daniel Hiltgen committed
92
93
94
	select {
	case s.pendingReqCh <- req:
	default:
95
		req.errCh <- ErrMaxQueue
Daniel Hiltgen's avatar
Daniel Hiltgen committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
	}
	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
120
			pending.schedAttempts++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
121
122
			if pending.origNumCtx == 0 {
				pending.origNumCtx = pending.opts.NumCtx
Daniel Hiltgen's avatar
Daniel Hiltgen committed
123
			}
124
125
126
127
128

			if pending.ctx.Err() != nil {
				slog.Debug("pending request cancelled or timed out, skipping scheduling")
				continue
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
129
130
131
132
133
134
135
136
137
			numParallel := envconfig.NumParallel
			// TODO (jmorganca): multimodal models don't support parallel yet
			// see https://github.com/ollama/ollama/issues/4165
			if len(pending.model.ProjectorPaths) > 0 && numParallel != 1 {
				numParallel = 1
				slog.Warn("multimodal models don't support parallel requests yet")
			}
			// Keep NumCtx and numParallel in sync
			if numParallel > 1 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
138
				pending.opts.NumCtx = pending.origNumCtx * numParallel
Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
			}
140

Daniel Hiltgen's avatar
Daniel Hiltgen committed
141
			for {
142
143
144
145
146
				cpus := s.getCpuFn()
				var systemMem gpu.GpuInfo
				if len(cpus) > 0 {
					systemMem = cpus[0]
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
147
148
149
150
151
152
153
154
155
156
157
158
159
				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) {
						runnerToExpire = runner
					} else {
						// Runner is usable, return it
						pending.useLoadedRunner(runner, s.finishedReqCh)
						break
					}
160
				} else if envconfig.MaxRunners > 0 && loadedCount >= envconfig.MaxRunners {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
161
					slog.Debug("max runners achieved, unloading one to make room", "runner_count", loadedCount)
162
					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
163
				} else {
164
					// Either no models are loaded or below envconfig.MaxRunners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
165
					// Get a refreshed GPU list
166
167
168
169
170
171
					var gpus gpu.GpuInfoList
					if pending.opts.NumGPU == 0 {
						gpus = s.getCpuFn()
					} else {
						gpus = s.getGpuFn()
					}
172

173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
					if envconfig.MaxRunners <= 0 {
						// 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 {
							envconfig.MaxRunners = defaultModelsPerGPU * len(gpus)
							slog.Debug("updating default concurrency", "OLLAMA_MAX_LOADED_MODELS", envconfig.MaxRunners, "gpu_count", len(gpus))
						} else {
							slog.Info("one or more GPUs detected that are unable to accurately report free memory - disabling default concurrency")
							envconfig.MaxRunners = len(gpus)
						}
					}
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
201
					estimate := llm.EstimateGPULayers(gpus, ggml, pending.model.ProjectorPaths, pending.opts)
					maxSize := systemMem.FreeMemory
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216

					// Add available GPU memory to the total pool
					// macOS hardware has unified memory so don't double count
					if runtime.GOOS != "darwin" {
						for _, gpu := range gpus {
							if gpu.Library == "cpu" {
								continue
							}
							if loadedCount == 0 {
								// If no other models are loaded, set the limit based on what's available
								maxSize += gpu.FreeMemory
							} else {
								// Other models could be unloaded, favor total memory for limit
								maxSize += gpu.TotalMemory
							}
217
218
						}
					}
219
220

					// Block attempting to load a model larger than system memory + GPU memory
221
222
					if estimate.TotalSize > maxSize {
						slog.Warn("model request too large for system", "requested", format.HumanBytes2(estimate.TotalSize), "system", format.HumanBytes2(maxSize))
223
224
225
226
227
228
229

						// Linux will crash if over-allocating memory - return an error to the user.
						// TODO (jmorganca): add reasonable upper limits for darwin and windows as well
						if runtime.GOOS == "linux" {
							pending.errCh <- fmt.Errorf("requested model (%s) is too large for this system (%s)", format.HumanBytes2(estimate.TotalSize), format.HumanBytes2(maxSize))
							break
						}
230
231
					}

232
233
					// 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
234
235
236
						// simplifying assumption of defaultParallel when in CPU mode
						if numParallel <= 0 {
							numParallel = defaultParallel
Daniel Hiltgen's avatar
Daniel Hiltgen committed
237
							pending.opts.NumCtx = pending.origNumCtx * numParallel
Daniel Hiltgen's avatar
Daniel Hiltgen committed
238
239
						}

240
241
						if loadedCount == 0 {
							slog.Debug("cpu mode with first model, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
242
							s.loadFn(pending, ggml, gpus, numParallel)
243
244
245
246
247
							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
248
							s.loadFn(pending, ggml, gpus, numParallel)
249
250
251
252
253
							break
						}
						// else we need to expire a runner
					} else if loadedCount == 0 {
						// No models loaded. Load the model but prefer the best fit.
254
						slog.Debug("loading first model", "model", pending.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
255
						g := pickBestFitGPUs(pending, ggml, gpus, &numParallel)
256
257
258
						if g != nil {
							gpus = g
						}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
259
						s.loadFn(pending, ggml, gpus, numParallel)
260
261
262
						break
					}

263
					if runnerToExpire == nil {
264
265
266
267
268
269
						// 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
270
						availGpus := s.filterGPUsWithoutLoadingModels(gpus)
271

272
						// Update free memory from currently loaded models
273
						s.updateFreeSpace(availGpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
274
						fitGpus := pickBestFitGPUs(pending, ggml, availGpus, &numParallel)
275
						if fitGpus != nil {
276
							slog.Debug("new model fits with existing models, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
277
							s.loadFn(pending, ggml, fitGpus, numParallel)
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
							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
							}()
297
298
299
							break
						}
						runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
300
301
302
303
304
305
306
307
308
309
					}
				}

				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()
310
				slog.Debug("resetting model to expire immediately to make room", "modelPath", runnerToExpire.modelPath, "refCount", runnerToExpire.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
311
312
313
314
315
316
317
318
319
320
321
322
				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.
323
				slog.Debug("waiting for pending requests to complete and unload to occur", "modelPath", runnerToExpire.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
324
325
326
327
328
				select {
				case <-ctx.Done():
					slog.Debug("shutting down scheduler pending loop")
					return
				case <-s.unloadedCh:
329
					slog.Debug("unload completed", "modelPath", runnerToExpire.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
					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 {
352
				slog.Error("finished request signal received after model unloaded", "modelPath", finished.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
353
354
355
356
357
358
				continue
			}
			runner.refMu.Lock()
			runner.refCount--
			if runner.refCount <= 0 {
				if runner.sessionDuration <= 0 {
359
					slog.Debug("runner with zero duration has gone idle, expiring to unload", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
360
361
362
363
364
365
					if runner.expireTimer != nil {
						runner.expireTimer.Stop()
						runner.expireTimer = nil
					}
					s.expiredCh <- runner
				} else if runner.expireTimer == nil {
366
					slog.Debug("runner with non-zero duration has gone idle, adding timer", "modelPath", runner.modelPath, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
367
					runner.expireTimer = time.AfterFunc(runner.sessionDuration, func() {
368
						slog.Debug("timer expired, expiring to unload", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
369
370
371
372
						runner.refMu.Lock()
						defer runner.refMu.Unlock()
						if runner.expireTimer != nil {
							runner.expireTimer.Stop()
373
							runner.expireTimer = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
374
375
376
						}
						s.expiredCh <- runner
					})
377
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
378
				} else {
379
					slog.Debug("runner with non-zero duration has gone idle, resetting timer", "modelPath", runner.modelPath, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
380
					runner.expireTimer.Reset(runner.sessionDuration)
381
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
382
383
				}
			}
384
			slog.Debug("after processing request finished event", "modelPath", runner.modelPath, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
385
386
			runner.refMu.Unlock()
		case runner := <-s.expiredCh:
387
			slog.Debug("runner expired event received", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
388
389
390
			runner.refMu.Lock()
			if runner.refCount > 0 {
				// Shouldn't happen, but safeguard to ensure no leaked runners
391
				slog.Debug("expired event with positive ref count, retrying", "modelPath", runner.modelPath, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
392
393
394
395
396
397
398
399
400
401
				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
			}

402
			s.loadedMu.Lock()
403
			slog.Debug("got lock to unload", "modelPath", runner.modelPath)
404
			finished := runner.waitForVRAMRecovery()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
405
			runner.unload()
406
			delete(s.loaded, runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
407
			s.loadedMu.Unlock()
408
			slog.Debug("runner released", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
409
			runner.refMu.Unlock()
410
411

			<-finished
412
			slog.Debug("sending an unloaded event", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
413
414
415
416
417
418
419
420
421
422
423
424
			s.unloadedCh <- struct{}{}
		}
	}
}

// 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++
425
426
427
428
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
429
430
431
	if pending.sessionDuration != nil {
		runner.sessionDuration = pending.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
432
433
434
435
436
437
438
439
	pending.successCh <- runner
	go func() {
		<-pending.ctx.Done()
		slog.Debug("context for request finished")
		finished <- pending
	}()
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
440
441
442
443
func (s *Scheduler) load(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList, numParallel int) {
	if numParallel < 1 {
		numParallel = 1
	}
444
445
446
447
	sessionDuration := envconfig.KeepAlive
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
448
	llama, err := s.newServerFn(gpus, req.model.ModelPath, ggml, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, numParallel)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
449
450
451
452
453
454
455
456
457
458
459
	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
		if errors.Is(llm.ErrUnsupportedFormat, err) || strings.Contains(err.Error(), "failed to load model") {
			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
	}
460
461
462
463
464
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
465
		sessionDuration: sessionDuration,
466
467
468
469
470
471
		gpus:            gpus,
		estimatedVRAM:   llama.EstimatedVRAM(),
		estimatedTotal:  llama.EstimatedTotal(),
		loading:         true,
		refCount:        1,
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
472
	runner.numParallel = numParallel
Daniel Hiltgen's avatar
Daniel Hiltgen committed
473
	runner.refMu.Lock()
474

Daniel Hiltgen's avatar
Daniel Hiltgen committed
475
476
477
478
479
480
481
482
483
484
485
	s.loadedMu.Lock()
	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)
			runner.refCount--
			req.errCh <- err
486
			slog.Debug("triggering expiration for failed load", "model", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
			s.expiredCh <- runner
			return
		}
		slog.Debug("finished setting up runner", "model", req.model.ModelPath)
		runner.loading = false
		go func() {
			<-req.ctx.Done()
			slog.Debug("context for request finished")
			s.finishedReqCh <- req
		}()
		req.successCh <- runner
	}()
}

func (s *Scheduler) updateFreeSpace(allGpus gpu.GpuInfoList) {
	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()
	for _, r := range s.loaded {
		r.refMu.Lock()
		if r.llama != nil {
			for _, gpu := range allGpus {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
512
				predMap[predKey{gpu.Library, gpu.ID}] += r.llama.EstimatedVRAMByGPU(gpu.ID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
			}
		} else {
			slog.Warn("unexpected nil runner reference, memory prediction may be incorrect")
		}
		r.refMu.Unlock()
	}
	s.loadedMu.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
				// after we start our first runner, then we'll never acount for that, so picking the smallest free value seems prudent.
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
535
536
537
538
539
540
541
542
543
			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)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
544
func (s *Scheduler) filterGPUsWithoutLoadingModels(allGpus gpu.GpuInfoList) gpu.GpuInfoList {
545
546
547
548
549
550
551
552
553
554
555
556
557
558
	ret := append(gpu.GpuInfoList{}, allGpus...)
	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
559
560
		}
	}
561
	return ret
Daniel Hiltgen's avatar
Daniel Hiltgen committed
562
563
}

564
// TODO consolidate sched_types.go
Daniel Hiltgen's avatar
Daniel Hiltgen committed
565
566
567
568
569
570
type runnerRef struct {
	refMu sync.Mutex
	// refCond   sync.Cond // Signaled on transition from 1 -> 0 refCount
	refCount uint // prevent unloading if > 0
	// unloading bool      // set to true when we are trying to unload the runner

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

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

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

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

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

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

616
617
618
619
620
621
622
623
	// 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
624
625
626
	// Normalize the NumCtx for parallelism
	optsExisting.NumCtx = optsExisting.NumCtx / runner.numParallel

627
	ctx, cancel := context.WithTimeout(ctx, timeout)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
628
	defer cancel()
629
630
	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
631
632
633
634
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
635

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

639
640
641
642
643
644
645
646
647
648
649
// 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
func (runner *runnerRef) waitForVRAMRecovery() chan interface{} {
	finished := make(chan interface{}, 1)

650
651
	// CPU or Metal don't need checking, so no waiting required
	// windows can page VRAM, only cuda currently can report accurate used vram usage
652
653
	if len(runner.gpus) == 0 ||
		(len(runner.gpus) == 1 && (runner.gpus[0].Library == "cpu" || runner.gpus[0].Library == "metal")) ||
654
		(runtime.GOOS == "windows" && runner.gpus[0].Library != "cuda") {
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
		finished <- struct{}{}
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
	gpusBefore := gpu.GetGPUInfo()
	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) {
674
				slog.Warn("gpu VRAM usage didn't recover within timeout", "seconds", time.Since(start).Seconds(), "model", runner.modelPath)
675
676
677
678
679
680
681
682
683
684
685
686
				finished <- struct{}{}
			}

			// Query GPUs, look for free to go back up
			gpusNow := gpu.GetGPUInfo()
			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 {
687
				slog.Debug(fmt.Sprintf("gpu VRAM free memory converged after %0.2f seconds", time.Since(start).Seconds()), "model", runner.modelPath)
688
689
690
691
692
693
694
695
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
type ByDuration []*runnerRef

func (a ByDuration) Len() int      { return len(a) }
func (a ByDuration) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a ByDuration) Less(i, j int) bool {
	// uint64 to turn negative time (never unload) to largest
	return uint64(a[i].sessionDuration) < uint64(a[j].sessionDuration)
}

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

// pickBestFitGPUs will try to find the optimal placement of the model in the available GPUs where the model fully fits
// If the model can not be fit fully within the available GPU(s) nil is returned
Daniel Hiltgen's avatar
Daniel Hiltgen committed
713
714
715
// If numParallel is <= 0, this will attempt try to optimize parallism based on available VRAM, and adjust
// opts.NumCtx accordingly
func pickBestFitGPUs(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList, numParallel *int) gpu.GpuInfoList {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
716
	var estimatedVRAM uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
717
718
719
720

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
726
727
728
729
730
	for _, gl := range gpus.ByLibrary() {
		var ok bool
		sgl := append(make(gpu.GpuInfoList, 0, len(gl)), gl...)

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

		// First attempt to fit the model into a single GPU
Daniel Hiltgen's avatar
Daniel Hiltgen committed
736
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
737
			req.opts.NumCtx = req.origNumCtx * p
Daniel Hiltgen's avatar
Daniel Hiltgen committed
738
739
740
741
742
743
744
			if !envconfig.SchedSpread {
				for _, g := range sgl {
					if ok, estimatedVRAM = llm.PredictServerFit([]gpu.GpuInfo{g}, ggml, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
						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
						return []gpu.GpuInfo{g}
					}
745
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
746
747
748
749
750
751
752
753
			}
		}

		// 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
754
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
755
			req.opts.NumCtx = req.origNumCtx * p
Daniel Hiltgen's avatar
Daniel Hiltgen committed
756
757
758
759
760
			if ok, estimatedVRAM = llm.PredictServerFit(sgl, ggml, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
				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
761
762
763
764
765
766
		}
	}
	return nil
}

// findRunnerToUnload finds a runner to unload to make room for a new model
767
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
768
769
770
771
772
773
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
774
775
776
777
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807

	// 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?
	sort.Sort(ByDuration(runnerList))

	// 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 {
			slog.Debug("found an idle runner to unload")
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
	slog.Debug("no idle runners, picking the shortest duration", "count", len(runnerList))
	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()
		}
	}
}
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822

// 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
func (s *Scheduler) maybeFindCPURunnerToUnload(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList) *runnerRef {
	slog.Debug("evaluating if CPU model load will fit in available system memory")
	estimate := llm.EstimateGPULayers(gpus, ggml, req.model.ProjectorPaths, req.opts)
	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()
}