sched.go 29 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"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
11
	"sort"
Michael Yang's avatar
int  
Michael Yang committed
12
	"strconv"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
13
14
15
16
17
	"strings"
	"sync"
	"time"

	"github.com/ollama/ollama/api"
18
	"github.com/ollama/ollama/discover"
19
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
20
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
21
	"github.com/ollama/ollama/fs/ggml"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
22
23
24
25
26
27
28
	"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
29
	origNumCtx      int // Track the initial ctx request
30
	sessionDuration *api.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
31
32
	successCh       chan *runnerRef
	errCh           chan error
33
	schedAttempts   uint
Daniel Hiltgen's avatar
Daniel Hiltgen committed
34
35
36
37
38
39
40
41
42
43
44
}

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

	loaded   map[string]*runnerRef
	loadedMu sync.Mutex

Michael Yang's avatar
Michael Yang committed
45
46
	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)
47
48
	getGpuFn     func() discover.GpuInfoList
	getCpuFn     func() discover.GpuInfoList
49
	reschedDelay time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
50
51
}

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

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

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

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

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

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

			if pending.ctx.Err() != nil {
				slog.Debug("pending request cancelled or timed out, skipping scheduling")
				continue
			}
Michael Yang's avatar
Michael Yang committed
133
			numParallel := int(envconfig.NumParallel())
134
			// TODO (jmorganca): mllama doesn't support parallel yet
Daniel Hiltgen's avatar
Daniel Hiltgen committed
135
			// see https://github.com/ollama/ollama/issues/4165
136
			if checkMllamaModelFamily(pending.model) && numParallel != 1 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
137
				numParallel = 1
138
				slog.Warn("mllama doesn't support parallel requests yet")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
			}
140

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

Michael Yang's avatar
int  
Michael Yang committed
168
					if envconfig.MaxRunners() <= 0 {
169
170
171
172
173
174
175
176
177
178
179
						// 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
180
181
							// HACK
							os.Setenv("OLLAMA_MAX_LOADED_MODELS", strconv.Itoa(defaultModelsPerGPU*len(gpus)))
182
							slog.Debug("updating default concurrency", "OLLAMA_MAX_LOADED_MODELS", envconfig.MaxRunners(), "gpu_count", len(gpus))
183
						} else {
Michael Yang's avatar
int  
Michael Yang committed
184
185
							// HACK
							os.Setenv("OLLAMA_MAX_LOADED_MODELS", strconv.Itoa(len(gpus)))
186
187
188
							slog.Info("one or more GPUs detected that are unable to accurately report free memory - disabling default concurrency")
						}
					}
189

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

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

202
203
					// 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
204
205
206
207
208
						// simplifying assumption of defaultParallel when in CPU mode
						if numParallel <= 0 {
							numParallel = defaultParallel
						}

209
210
						pending.opts.NumCtx = pending.origNumCtx * numParallel

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

237
					if runnerToExpire == nil {
238
239
240
241
242
243
						// 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
244
						availGpus := s.filterGPUsWithoutLoadingModels(gpus)
245

246
						// Update free memory from currently loaded models
247
						s.updateFreeSpace(availGpus)
248
						fitGpus := pickBestFullFitByLibrary(pending, ggml, availGpus, &numParallel)
249
						if fitGpus != nil {
250
							slog.Debug("new model fits with existing models, loading")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
251
							s.loadFn(pending, ggml, fitGpus, numParallel)
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
							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
							}()
271
272
273
							break
						}
						runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
274
275
276
277
278
279
280
281
282
283
					}
				}

				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()
284
				slog.Debug("resetting model to expire immediately to make room", "modelPath", runnerToExpire.modelPath, "refCount", runnerToExpire.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
285
286
287
288
289
290
291
292
293
294
295
296
				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.
297
				slog.Debug("waiting for pending requests to complete and unload to occur", "modelPath", runnerToExpire.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
298
299
300
301
302
				select {
				case <-ctx.Done():
					slog.Debug("shutting down scheduler pending loop")
					return
				case <-s.unloadedCh:
303
					slog.Debug("unload completed", "modelPath", runnerToExpire.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
					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 {
326
				slog.Error("finished request signal received after model unloaded", "modelPath", finished.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
327
328
329
330
331
332
				continue
			}
			runner.refMu.Lock()
			runner.refCount--
			if runner.refCount <= 0 {
				if runner.sessionDuration <= 0 {
333
					slog.Debug("runner with zero duration has gone idle, expiring to unload", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
334
335
336
337
338
339
					if runner.expireTimer != nil {
						runner.expireTimer.Stop()
						runner.expireTimer = nil
					}
					s.expiredCh <- runner
				} else if runner.expireTimer == nil {
340
					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
341
					runner.expireTimer = time.AfterFunc(runner.sessionDuration, func() {
342
						slog.Debug("timer expired, expiring to unload", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
343
344
345
346
						runner.refMu.Lock()
						defer runner.refMu.Unlock()
						if runner.expireTimer != nil {
							runner.expireTimer.Stop()
347
							runner.expireTimer = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
348
349
350
						}
						s.expiredCh <- runner
					})
351
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
352
				} else {
353
					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
354
					runner.expireTimer.Reset(runner.sessionDuration)
355
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
356
357
				}
			}
358
			slog.Debug("after processing request finished event", "modelPath", runner.modelPath, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
359
360
			runner.refMu.Unlock()
		case runner := <-s.expiredCh:
361
			slog.Debug("runner expired event received", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
362
363
			runner.refMu.Lock()
			if runner.refCount > 0 {
364
				slog.Debug("expired event with positive ref count, retrying", "modelPath", runner.modelPath, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
365
366
367
368
369
370
371
372
373
374
				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
			}

375
			s.loadedMu.Lock()
376
			slog.Debug("got lock to unload", "modelPath", runner.modelPath)
377
			finished := runner.waitForVRAMRecovery()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
378
			runner.unload()
379
			delete(s.loaded, runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
380
			s.loadedMu.Unlock()
381
			slog.Debug("runner released", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
382
			runner.refMu.Unlock()
383
384

			<-finished
385
			slog.Debug("sending an unloaded event", "modelPath", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
386
387
388
389
390
391
392
393
394
395
396
397
			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++
398
399
400
401
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
402
403
404
	if pending.sessionDuration != nil {
		runner.sessionDuration = pending.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
405
406
407
408
409
410
411
412
	pending.successCh <- runner
	go func() {
		<-pending.ctx.Done()
		slog.Debug("context for request finished")
		finished <- pending
	}()
}

Michael Yang's avatar
Michael Yang committed
413
func (s *Scheduler) load(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel int) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
414
415
416
	if numParallel < 1 {
		numParallel = 1
	}
Michael Yang's avatar
Michael Yang committed
417
	sessionDuration := envconfig.KeepAlive()
418
419
420
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Michael Yang's avatar
Michael Yang committed
421
	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
422
423
424
425
	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
426
		if errors.Is(err, ggml.ErrUnsupportedFormat) || strings.Contains(err.Error(), "failed to load model") {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
427
428
429
430
431
432
			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
	}
433
434
435
436
437
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
438
		sessionDuration: sessionDuration,
439
440
441
442
443
444
		gpus:            gpus,
		estimatedVRAM:   llama.EstimatedVRAM(),
		estimatedTotal:  llama.EstimatedTotal(),
		loading:         true,
		refCount:        1,
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
445
	runner.numParallel = numParallel
Daniel Hiltgen's avatar
Daniel Hiltgen committed
446
	runner.refMu.Lock()
447

Daniel Hiltgen's avatar
Daniel Hiltgen committed
448
449
450
451
452
453
454
455
456
457
458
	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
459
			slog.Debug("triggering expiration for failed load", "model", runner.modelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
460
461
462
463
464
465
466
467
468
469
470
471
472
473
			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
	}()
}

474
func (s *Scheduler) updateFreeSpace(allGpus discover.GpuInfoList) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
475
476
477
478
479
480
481
482
483
484
	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
485
				predMap[predKey{gpu.Library, gpu.ID}] += r.llama.EstimatedVRAMByGPU(gpu.ID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
			}
		} 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
505
				// 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
506
507
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
508
509
510
511
512
513
514
515
516
			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)
517
518
func (s *Scheduler) filterGPUsWithoutLoadingModels(allGpus discover.GpuInfoList) discover.GpuInfoList {
	ret := append(discover.GpuInfoList{}, allGpus...)
519
520
521
522
523
524
525
526
527
528
529
530
531
	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
532
533
		}
	}
534
	return ret
Daniel Hiltgen's avatar
Daniel Hiltgen committed
535
536
}

537
// TODO consolidate sched_types.go
Daniel Hiltgen's avatar
Daniel Hiltgen committed
538
539
540
541
542
543
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

544
	llama          llm.LlamaServer
545
546
	loading        bool                 // True only during initial load, then false forever
	gpus           discover.GpuInfoList // Recorded at time of provisioning
547
548
	estimatedVRAM  uint64
	estimatedTotal uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
549
550
551

	sessionDuration time.Duration
	expireTimer     *time.Timer
552
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
553

Daniel Hiltgen's avatar
Daniel Hiltgen committed
554
555
556
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
557
558
559
560
561
	*api.Options
}

// The refMu must already be held when calling unload
func (runner *runnerRef) unload() {
562
563
564
565
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
566
567
568
	if runner.llama != nil {
		runner.llama.Close()
	}
569
	runner.model = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
570
571
572
573
574
575
576
577
578
	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()
579

Daniel Hiltgen's avatar
Daniel Hiltgen committed
580
581
582
583
	timeout := 10 * time.Second
	if runner.loading {
		timeout = 2 * time.Minute // Initial load can take a long time for big models on slow systems...
	}
584

585
586
587
588
	if runner.Options == nil {
		return true
	}

589
590
591
592
593
594
595
596
	// 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
597
598
599
	// Normalize the NumCtx for parallelism
	optsExisting.NumCtx = optsExisting.NumCtx / runner.numParallel

600
	ctx, cancel := context.WithTimeout(ctx, timeout)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
601
	defer cancel()
602
603
	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
604
605
606
607
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
608

Daniel Hiltgen's avatar
Daniel Hiltgen committed
609
610
611
	return false
}

612
613
614
615
616
617
618
619
620
621
622
// 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)

623
624
	// CPU or Metal don't need checking, so no waiting required
	// windows can page VRAM, only cuda currently can report accurate used vram usage
625
626
	if len(runner.gpus) == 0 ||
		(len(runner.gpus) == 1 && (runner.gpus[0].Library == "cpu" || runner.gpus[0].Library == "metal")) ||
627
		(runtime.GOOS == "windows" && runner.gpus[0].Library != "cuda") {
628
629
630
631
632
633
		finished <- struct{}{}
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
634
	gpusBefore := discover.GetGPUInfo()
635
636
637
638
639
640
641
642
643
644
645
646
	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) {
647
				slog.Warn("gpu VRAM usage didn't recover within timeout", "seconds", time.Since(start).Seconds(), "model", runner.modelPath)
648
649
650
651
				finished <- struct{}{}
			}

			// Query GPUs, look for free to go back up
652
			gpusNow := discover.GetGPUInfo()
653
654
655
656
657
658
659
			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 {
660
				slog.Debug(fmt.Sprintf("gpu VRAM free memory converged after %0.2f seconds", time.Since(start).Seconds()), "model", runner.modelPath)
661
662
663
664
665
666
667
668
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
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 }

684
685
// 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
686
// If the model can not be fit fully within the available GPU(s) nil is returned
687
// If numParallel is <= 0, this will attempt try to optimize parallelism based on available VRAM, and adjust
Daniel Hiltgen's avatar
Daniel Hiltgen committed
688
// opts.NumCtx accordingly
Michael Yang's avatar
Michael Yang committed
689
func pickBestFullFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel *int) discover.GpuInfoList {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
690
	var estimatedVRAM uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
691
692
693
694

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
700
701
	for _, gl := range gpus.ByLibrary() {
		var ok bool
702
		sgl := append(make(discover.GpuInfoList, 0, len(gl)), gl...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
703
704

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

		// First attempt to fit the model into a single GPU
Daniel Hiltgen's avatar
Daniel Hiltgen committed
710
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
711
			req.opts.NumCtx = req.origNumCtx * p
Michael Yang's avatar
bool  
Michael Yang committed
712
			if !envconfig.SchedSpread() {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
713
				for _, g := range sgl {
Michael Yang's avatar
Michael Yang committed
714
					if ok, estimatedVRAM = llm.PredictServerFit([]discover.GpuInfo{g}, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
715
716
						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
717
						return []discover.GpuInfo{g}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
718
					}
719
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
720
721
722
723
724
725
726
727
			}
		}

		// 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
728
		for _, p := range numParallelToTry {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
729
			req.opts.NumCtx = req.origNumCtx * p
Michael Yang's avatar
Michael Yang committed
730
			if ok, estimatedVRAM = llm.PredictServerFit(sgl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
731
732
733
734
				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
735
736
737
738
739
		}
	}
	return nil
}

740
// If multiple Libraries are detected, pick the Library which loads the most layers for the model
Michael Yang's avatar
Michael Yang committed
741
func pickBestPartialFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, numParallel *int) discover.GpuInfoList {
742
743
	if *numParallel <= 0 {
		*numParallel = 1
Richard Lyons's avatar
Richard Lyons committed
744
745
		req.opts.NumCtx = req.origNumCtx
	}
746
747
748
749
750
751
752
	byLibrary := gpus.ByLibrary()
	if len(byLibrary) <= 1 {
		return gpus
	}
	var bestEstimate uint64
	var bestFit int
	for i, gl := range byLibrary {
Michael Yang's avatar
Michael Yang committed
753
		_, estimatedVRAM := llm.PredictServerFit(gl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts)
754
755
756
757
758
759
760
761
		if estimatedVRAM > bestEstimate {
			bestEstimate = estimatedVRAM
			bestFit = i
		}
	}
	return byLibrary[bestFit]
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
762
// findRunnerToUnload finds a runner to unload to make room for a new model
763
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
764
765
766
767
768
769
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
770
771
772
773
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
774
775
776
777
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

	// 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()
		}
	}
}
804

Patrick Devine's avatar
Patrick Devine committed
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
func (s *Scheduler) expireRunner(model *Model) {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
	runner, ok := s.loaded[model.ModelPath]
	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()
	}
}

824
825
// 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
826
func (s *Scheduler) maybeFindCPURunnerToUnload(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList) *runnerRef {
827
	slog.Debug("evaluating if CPU model load will fit in available system memory")
Michael Yang's avatar
Michael Yang committed
828
	estimate := llm.EstimateGPULayers(gpus, f, req.model.ProjectorPaths, req.opts)
829
830
831
832
833
834
835
836
837
	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()
}