sched.go 29 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
	"slices"
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/discover"
17
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
18
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
19
	"github.com/ollama/ollama/fs/ggml"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
20
	"github.com/ollama/ollama/llm"
21
	"github.com/ollama/ollama/logutil"
22
	"github.com/ollama/ollama/ml"
23
	"github.com/ollama/ollama/types/model"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
24
25
26
27
28
29
)

type LlmRequest struct {
	ctx             context.Context //nolint:containedctx
	model           *Model
	opts            api.Options
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
}

type Scheduler struct {
	pendingReqCh  chan *LlmRequest
	finishedReqCh chan *LlmRequest
	expiredCh     chan *runnerRef
40
	unloadedCh    chan any
Daniel Hiltgen's avatar
Daniel Hiltgen committed
41

Jesse Gross's avatar
Jesse Gross committed
42
	// loadedMu protects loaded and activeLoading
Daniel Hiltgen's avatar
Daniel Hiltgen committed
43
44
	loadedMu sync.Mutex

Jesse Gross's avatar
Jesse Gross committed
45
46
47
48
49
50
51
	// activeLoading is the model that we are currently working on loading,
	// including by evicting one or more other models. We can only load
	// one model at a time but new requests to models that already loaded can
	// happen in parallel
	activeLoading llm.LlamaServer
	loaded        map[string]*runnerRef

52
53
54
55
	loadFn          func(req *LlmRequest, f *ggml.GGML, systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) bool
	newServerFn     func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int) (llm.LlamaServer, error)
	getGpuFn        func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo
	getSystemInfoFn func() ml.SystemInfo
56
	waitForRecovery time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
57
58
}

59
60
61
62
63
// 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

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{
69
70
71
72
73
74
		pendingReqCh:    make(chan *LlmRequest, maxQueue),
		finishedReqCh:   make(chan *LlmRequest, maxQueue),
		expiredCh:       make(chan *runnerRef, maxQueue),
		unloadedCh:      make(chan any, maxQueue),
		loaded:          make(map[string]*runnerRef),
		newServerFn:     llm.NewLlamaServer,
75
76
		getGpuFn:        discover.GPUDevices,
		getSystemInfoFn: discover.GetSystemInfo,
77
		waitForRecovery: 5 * time.Second,
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
Jesse Gross's avatar
Jesse Gross committed
84
func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, sessionDuration *api.Duration) (chan *runnerRef, chan error) {
85
86
87
88
	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

Jesse Gross's avatar
Jesse Gross committed
89
90
91
92
93
	if m.CheckCapabilities(model.CapabilityVision) == nil {
		// multimodal models require at least 2048 context
		opts.NumCtx = max(opts.NumCtx, 2048)
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
94
95
	req := &LlmRequest{
		ctx:             c,
Jesse Gross's avatar
Jesse Gross committed
96
		model:           m,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
97
98
		opts:            opts,
		sessionDuration: sessionDuration,
Jesse Gross's avatar
Jesse Gross committed
99
		successCh:       make(chan *runnerRef, 1),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
100
101
		errCh:           make(chan error, 1),
	}
102

Jesse Gross's avatar
Jesse Gross committed
103
104
105
106
107
108
109
110
111
112
113
	s.loadedMu.Lock()
	runner := s.loaded[req.model.ModelPath]
	s.loadedMu.Unlock()
	if runner != nil && !runner.needsReload(c, req) {
		req.useLoadedRunner(runner, s.finishedReqCh)
	} else {
		select {
		case s.pendingReqCh <- req:
		default:
			req.errCh <- ErrMaxQueue
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
	}
	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) {
131
132
	maxRunners := envconfig.MaxRunners()

Daniel Hiltgen's avatar
Daniel Hiltgen committed
133
134
135
136
137
138
139
	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
140
			pending.schedAttempts++
141
142
143
144
145

			if pending.ctx.Err() != nil {
				slog.Debug("pending request cancelled or timed out, skipping scheduling")
				continue
			}
146
			logutil.Trace("processing incoming request", "model", pending.model.ModelPath)
147

Daniel Hiltgen's avatar
Daniel Hiltgen committed
148
149
150
151
152
			for {
				var runnerToExpire *runnerRef
				s.loadedMu.Lock()
				runner := s.loaded[pending.model.ModelPath]
				loadedCount := len(s.loaded)
153
				runnersSnapshot := make([]ml.FilteredRunnerDiscovery, 0, len(s.loaded))
154
155
156
				for _, r := range s.loaded {
					runnersSnapshot = append(runnersSnapshot, r)
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
157
				s.loadedMu.Unlock()
158

Daniel Hiltgen's avatar
Daniel Hiltgen committed
159
160
				if runner != nil {
					if runner.needsReload(ctx, pending) {
161
						slog.Debug("reloading", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
162
163
164
						runnerToExpire = runner
					} else {
						// Runner is usable, return it
165
						logutil.Trace("using existing loaded runner", "model", pending.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
166
167
168
						pending.useLoadedRunner(runner, s.finishedReqCh)
						break
					}
169
				} else if maxRunners > 0 && loadedCount >= int(maxRunners) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
170
					slog.Debug("max runners achieved, unloading one to make room", "runner_count", loadedCount)
171
					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
172
				} else {
173
					// Either no models are loaded or below envconfig.MaxRunners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
174
					// Get a refreshed GPU list
175
					var gpus []ml.DeviceInfo
176
					if pending.opts.NumGPU == 0 {
177
						gpus = []ml.DeviceInfo{}
178
					} else {
179
						logutil.Trace("refreshing GPU list", "model", pending.model.ModelPath)
180
						gpus = s.getGpuFn(ctx, runnersSnapshot)
181
					}
182
					logutil.Trace("refreshing system information", "model", pending.model.ModelPath)
183
184
185
186
187
					systemInfo := s.getSystemInfoFn()
					if maxRunners <= 0 {
						// No user specified MaxRunners, so figure out what automatic setting to use for the next load attempt
						if pending.opts.NumGPU == 0 {
							// Need to get actual GPU list to set the correct default max models
188
							logutil.Trace("refreshing GPU list", "model", pending.model.ModelPath)
189
190
							g := s.getGpuFn(ctx, runnersSnapshot)
							maxRunners = uint(defaultModelsPerGPU * max(len(g), 1))
191
						} else {
192
							maxRunners = uint(defaultModelsPerGPU * max(len(gpus), 1))
193
						}
194
						slog.Debug("updating default concurrency", "OLLAMA_MAX_LOADED_MODELS", maxRunners, "gpu_count", len(gpus))
195
					}
196

197
					// Load model for fitting
198
					logutil.Trace("loading model metadata", "model", pending.model.ModelPath)
199
					ggml, err := llm.LoadModel(pending.model.ModelPath, 1024)
200
201
202
203
					if err != nil {
						pending.errCh <- err
						break
					}
204

Jesse Gross's avatar
Jesse Gross committed
205
					// Update free memory from currently loaded models
206
					logutil.Trace("updating free space", "gpu_count", len(gpus), "model", pending.model.ModelPath)
Jesse Gross's avatar
Jesse Gross committed
207
					s.updateFreeSpace(gpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
208

Jesse Gross's avatar
Jesse Gross committed
209
					if loadedCount == 0 {
210
						// No models loaded. Load the model but prefer the best fit.
211
						slog.Debug("loading first model", "model", pending.model.ModelPath)
212
						s.loadFn(pending, ggml, systemInfo, gpus, false)
213
214
215
						break
					}

Jesse Gross's avatar
Jesse Gross committed
216
217
					// More than one loaded model, so we have to see if the
					// new one fits
218
					logutil.Trace("loading additional model", "model", pending.model.ModelPath)
219
					needEvict := s.loadFn(pending, ggml, systemInfo, gpus, true)
Jesse Gross's avatar
Jesse Gross committed
220
221
222
					if !needEvict {
						slog.Debug("new model fits with existing models, loading")
						break
Daniel Hiltgen's avatar
Daniel Hiltgen committed
223
					}
Jesse Gross's avatar
Jesse Gross committed
224
225

					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
226
227
228
				}

				if runnerToExpire == nil {
229
230
231
					// While we were performing load calculations, the loaded runner(s) unloaded in parallel
					// so findRunnerToUnload returned no runners.  We'll try again and the loadedCount should be zero
					slog.Debug("runner to expire was nil, retrying")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
232
233
234
235
					continue
				}
				// Trigger an expiration to unload once it's done
				runnerToExpire.refMu.Lock()
236
				slog.Debug("resetting model to expire immediately to make room", "runner", runnerToExpire, "refCount", runnerToExpire.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
237
238
239
240
241
242
243
244
245
246
				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
247
				slog.Debug("waiting for pending requests to complete and unload to occur", "runner", runnerToExpire)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
248
249
250
251
252
				select {
				case <-ctx.Done():
					slog.Debug("shutting down scheduler pending loop")
					return
				case <-s.unloadedCh:
253
					slog.Debug("unload completed", "runner", runnerToExpire)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
					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 {
276
				slog.Error("finished request signal received after model unloaded", "modelPath", finished.model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
277
278
279
280
281
282
				continue
			}
			runner.refMu.Lock()
			runner.refCount--
			if runner.refCount <= 0 {
				if runner.sessionDuration <= 0 {
283
					slog.Debug("runner with zero duration has gone idle, expiring to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
284
285
286
287
288
289
					if runner.expireTimer != nil {
						runner.expireTimer.Stop()
						runner.expireTimer = nil
					}
					s.expiredCh <- runner
				} else if runner.expireTimer == nil {
290
					slog.Debug("runner with non-zero duration has gone idle, adding timer", "runner", runner, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
291
					runner.expireTimer = time.AfterFunc(runner.sessionDuration, func() {
292
						slog.Debug("timer expired, expiring to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
293
294
295
296
						runner.refMu.Lock()
						defer runner.refMu.Unlock()
						if runner.expireTimer != nil {
							runner.expireTimer.Stop()
297
							runner.expireTimer = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
298
299
300
						}
						s.expiredCh <- runner
					})
301
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
302
				} else {
303
					slog.Debug("runner with non-zero duration has gone idle, resetting timer", "runner", runner, "duration", runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
304
					runner.expireTimer.Reset(runner.sessionDuration)
305
					runner.expiresAt = time.Now().Add(runner.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
306
307
				}
			}
308
			slog.Debug("after processing request finished event", "runner", runner, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
309
310
			runner.refMu.Unlock()
		case runner := <-s.expiredCh:
311
			slog.Debug("runner expired event received", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
312
313
			runner.refMu.Lock()
			if runner.refCount > 0 {
314
				slog.Debug("expired event with positive ref count, retrying", "runner", runner, "refCount", runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
315
316
317
318
319
320
321
322
323
324
				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
			}

325
			s.loadedMu.Lock()
326
327
328
329
330
331
332
333
334
335
336
			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)
337
338
339
340
341
342
343
344
345
346
347
			} else if runner.pid != runnerToUnload.pid {
				// If the pids do not match, we likely had multiple load
				// failures for the same model in quick succession due to
				// request context canceled and are draining the queue of
				// events. Ensure the orphaned runner is properly shut down, but
				// do not delete the mismatched loaded runner, or wait for VRAM
				// convergence.
				slog.Debug("orphaned runner shutting down", "orphan", runner, "loaded", runnerToUnload)
				runner.unload()
				s.loadedMu.Unlock()
				runner.refMu.Unlock()
348
349
			} else {
				slog.Debug("starting background wait for VRAM recovery", "runner", runner)
350
				runnersSnapshot := make([]ml.FilteredRunnerDiscovery, 0, len(s.loaded))
351
352
353
354
				for _, r := range s.loaded {
					runnersSnapshot = append(runnersSnapshot, r)
				}
				finished := s.waitForVRAMRecovery(runner, runnersSnapshot)
355
356
357
358
359
360
361
362
363
				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
364
365
366
367
368
369
370
371
372
373
374
		}
	}
}

// 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++
375
376
377
378
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
379
380
381
	if pending.sessionDuration != nil {
		runner.sessionDuration = pending.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
382
383
384
	pending.successCh <- runner
	go func() {
		<-pending.ctx.Done()
385
		slog.Debug("context for request finished", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
386
387
388
389
		finished <- pending
	}()
}

Jesse Gross's avatar
Jesse Gross committed
390
391
// load creates a new model based on req and loads it. If requireFull is true then the model must be loaded fully onto GPUs
// (if any). Returns whether the scheduler needs to evict a model to make this one fit.
392
func (s *Scheduler) load(req *LlmRequest, f *ggml.GGML, systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) bool {
393
	numParallel := max(int(envconfig.NumParallel()), 1)
Jesse Gross's avatar
Jesse Gross committed
394
395
396
397
398
399

	// Embedding models should always be loaded with parallel=1
	if req.model.CheckCapabilities(model.CapabilityCompletion) != nil {
		numParallel = 1
	}

400
	// `mllama`, `qwen3vl`, and `qwen3vlmoe` are snowflakes and uses an encoder cache which cannot be used with num_parallel > 1
Jesse Gross's avatar
Jesse Gross committed
401
	// ref: https://github.com/ollama/ollama/issues/4165
402
	if slices.Contains([]string{"mllama", "qwen3vl", "qwen3vlmoe"}, req.model.Config.ModelFamily) && numParallel != 1 {
Jesse Gross's avatar
Jesse Gross committed
403
		numParallel = 1
404
		slog.Warn("model architecture does not currently support parallel requests", "architecture", req.model.Config.ModelFamily)
Jesse Gross's avatar
Jesse Gross committed
405
406
	}

Michael Yang's avatar
Michael Yang committed
407
	sessionDuration := envconfig.KeepAlive()
408
409
410
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Jesse Gross's avatar
Jesse Gross committed
411
412
413
414
415
416

	s.loadedMu.Lock()
	llama := s.activeLoading

	if llama == nil {
		var err error
417
		llama, err = s.newServerFn(systemInfo, gpus, req.model.ModelPath, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, numParallel)
Jesse Gross's avatar
Jesse Gross committed
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
		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(err, ggml.ErrUnsupportedFormat) || 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
			s.loadedMu.Unlock()
			return false
		}

		s.activeLoading = llama
	} else {
		if s.activeLoading.ModelPath() != req.model.ModelPath {
			panic(fmt.Errorf("attempting to load different model after eviction (original %v new %v)", s.activeLoading.ModelPath(), req.model.ModelPath))
		}
	}

	s.loadedMu.Unlock()

440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
	systemTotalMemory := systemInfo.TotalMemory
	systemFreeMemory := systemInfo.FreeMemory
	systemSwapFreeMemory := systemInfo.FreeSwap
	slog.Info("system memory", "total", format.HumanBytes2(systemTotalMemory), "free", format.HumanBytes2(systemFreeMemory), "free_swap", format.HumanBytes2(systemSwapFreeMemory))

	for _, gpu := range gpus {
		available := gpu.FreeMemory - envconfig.GpuOverhead() - gpu.MinimumMemory()
		if gpu.FreeMemory < envconfig.GpuOverhead()+gpu.MinimumMemory() {
			available = 0
		}
		slog.Info("gpu memory", "id", gpu.ID, "library", gpu.Library,
			"available", format.HumanBytes2(available),
			"free", format.HumanBytes2(gpu.FreeMemory),
			"minimum", format.HumanBytes2(gpu.MinimumMemory()),
			"overhead", format.HumanBytes2(envconfig.GpuOverhead()))
	}

457
	gpuIDs, err := llama.Load(req.ctx, systemInfo, gpus, requireFull)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
458
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
459
		if errors.Is(err, llm.ErrLoadRequiredFull) {
460
461
462
463
464
465
466
			if !requireFull {
				// No other models loaded, yet we still don't fit, so report an error
				slog.Info("model is too large for system memory", "requireFull", requireFull)
				s.activeLoading.Close()
				s.activeLoading = nil
				req.errCh <- err
			}
Jesse Gross's avatar
Jesse Gross committed
467
			return true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
468
		}
Jesse Gross's avatar
Jesse Gross committed
469
470
471
472

		slog.Info("Load failed", "model", req.model.ModelPath, "error", err)
		s.activeLoading.Close()
		s.activeLoading = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
473
		req.errCh <- err
Jesse Gross's avatar
Jesse Gross committed
474
		return false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
475
	}
Jesse Gross's avatar
Jesse Gross committed
476

477
478
479
480
481
482
483
484
485
486
487
488
489
490
	// Determine if we have discrete GPUs which we should monitor VRAM usage on during shutdown
	discreteGPUs := false
iGPUScan:
	for _, devid := range gpuIDs {
		for _, dev := range gpus {
			if dev.DeviceID == devid {
				if !dev.Integrated {
					discreteGPUs = true
					break iGPUScan
				}
			}
		}
	}

491
492
493
494
495
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
496
		sessionDuration: sessionDuration,
497
		gpus:            gpuIDs,
498
		discreteGPUs:    discreteGPUs,
Jesse Gross's avatar
Jesse Gross committed
499
500
		vramSize:        llama.VRAMSize(),
		totalSize:       llama.TotalSize(),
501
		loading:         true,
502
		pid:             llama.Pid(),
503
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
504
	runner.numParallel = numParallel
505
	runner.refMu.Lock() // hold lock until running or aborted
506

Daniel Hiltgen's avatar
Daniel Hiltgen committed
507
	s.loadedMu.Lock()
508
509
510
511
512
513
514
	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()
	}
Jesse Gross's avatar
Jesse Gross committed
515
	s.activeLoading = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
516
517
518
519
520
521
522
523
524
	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
525
			slog.Debug("triggering expiration for failed load", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
526
527
528
			s.expiredCh <- runner
			return
		}
529
		slog.Debug("finished setting up", "runner", runner)
530
531
532
		if runner.pid < 0 {
			runner.pid = llama.Pid()
		}
533
		runner.refCount++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
534
535
536
537
538
539
540
541
		runner.loading = false
		go func() {
			<-req.ctx.Done()
			slog.Debug("context for request finished")
			s.finishedReqCh <- req
		}()
		req.successCh <- runner
	}()
Jesse Gross's avatar
Jesse Gross committed
542
543

	return false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
544
545
}

546
547
548
549
func (s *Scheduler) updateFreeSpace(allGpus []ml.DeviceInfo) {
	if len(allGpus) == 0 {
		return
	}
550
	predMap := map[ml.DeviceID]uint64{} // Sum up the total predicted usage per GPU for all runners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
551
	s.loadedMu.Lock()
552
	runners := make([]*runnerRef, 0, len(s.loaded))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
553
	for _, r := range s.loaded {
554
555
556
557
		runners = append(runners, r)
	}
	s.loadedMu.Unlock()
	for _, r := range runners {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
558
559
560
		r.refMu.Lock()
		if r.llama != nil {
			for _, gpu := range allGpus {
561
				predMap[gpu.DeviceID] += r.llama.VRAMByGPU(gpu.DeviceID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
562
563
564
565
566
567
568
569
570
			}
		} 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 {
571
		if p, ok := predMap[allGpus[i].DeviceID]; ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
572
573
574
575
576
577
578
579
			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
580
				// 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
581
582
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
583
584
585
586
587
588
			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))
		}
	}
}

// TODO consolidate sched_types.go
Daniel Hiltgen's avatar
Daniel Hiltgen committed
589
type runnerRef struct {
590
	refMu    sync.Mutex
Daniel Hiltgen's avatar
Daniel Hiltgen committed
591
592
	refCount uint // prevent unloading if > 0

593
594
595
596
597
598
599
	llama        llm.LlamaServer
	pid          int
	loading      bool          // True only during initial load, then false forever
	gpus         []ml.DeviceID // Recorded at time of provisioning
	discreteGPUs bool          // True if all devices are discrete GPUs - used to skip VRAM recovery check for iGPUs
	vramSize     uint64
	totalSize    uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
600
601
602

	sessionDuration time.Duration
	expireTimer     *time.Timer
603
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
604

Daniel Hiltgen's avatar
Daniel Hiltgen committed
605
606
607
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
608
609
610
611
612
	*api.Options
}

// The refMu must already be held when calling unload
func (runner *runnerRef) unload() {
613
614
615
616
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
617
618
619
	if runner.llama != nil {
		runner.llama.Close()
	}
620
	runner.model = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
621
622
623
624
625
626
627
628
	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()
629

Daniel Hiltgen's avatar
Daniel Hiltgen committed
630
631
632
633
	timeout := 10 * time.Second
	if runner.loading {
		timeout = 2 * time.Minute // Initial load can take a long time for big models on slow systems...
	}
634

635
636
637
638
	if runner.Options == nil {
		return true
	}

639
640
641
642
643
644
645
646
647
	// 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
	}

	ctx, cancel := context.WithTimeout(ctx, timeout)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
648
	defer cancel()
649
650
	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
651
652
653
654
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
655

Daniel Hiltgen's avatar
Daniel Hiltgen committed
656
657
658
	return false
}

659
660
661
662
663
664
665
666
// 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
667
func (s *Scheduler) waitForVRAMRecovery(runner *runnerRef, runners []ml.FilteredRunnerDiscovery) chan any {
668
	finished := make(chan any, 1)
669

670
671
672
	// CPU, Metal and iGPUs don't need checking, so no waiting required
	if len(runner.gpus) == 0 || !runner.discreteGPUs ||
		(len(runner.gpus) == 1 && runner.gpus[0].Library == "Metal") {
673
		finished <- struct{}{}
674
		slog.Debug("no need to wait for VRAM recovery", "runner", runner)
675
676
677
678
679
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
680
	gpusBefore := s.getGpuFn(context.Background(), runners)
681
682
683
684
685
	var totalMemoryBefore, freeMemoryBefore uint64
	for _, gpu := range gpusBefore {
		totalMemoryBefore += gpu.TotalMemory
		freeMemoryBefore += gpu.FreeMemory
	}
686
687
688
	totalMemoryNow := totalMemoryBefore
	freeMemoryNow := freeMemoryBefore

689
	go func() {
690
691
		// typical convergence is 0.5-1.5s - If it takes too long to discover and converge, let the scheduler estimate VRAM usage
		ctx, cancel := context.WithTimeout(context.Background(), s.waitForRecovery)
692
		defer cancel()
693
694
695
		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
696
697
698
699
700
701
702
703
704
705
			select {
			case <-ticker.C:
				// Query GPUs, look for free to go back up
				gpusNow := s.getGpuFn(ctx, runners)
				totalMemoryNow = 0
				freeMemoryNow = 0
				for _, gpu := range gpusNow {
					totalMemoryNow += gpu.TotalMemory
					freeMemoryNow += gpu.FreeMemory
				}
706
707
708
709
710
				if freeMemoryNow > freeMemoryBefore {
					logutil.Trace("gpu VRAM convergence", "percent", int(float32(freeMemoryNow-freeMemoryBefore)/float32(runner.vramSize)*100))
				} else {
					logutil.Trace("gpu VRAM convergence", "percent", 0)
				}
711
712
713
714
715
716
717
718
				// If we're within ~75% of the estimated memory usage recovered, bail out
				if float32(freeMemoryNow-freeMemoryBefore) > float32(runner.vramSize)*0.75 {
					slog.Debug(fmt.Sprintf("gpu VRAM free memory converged after %0.2f seconds", time.Since(start).Seconds()), "free_before", format.HumanBytes2(freeMemoryBefore), "free_now", format.HumanBytes2(freeMemoryNow), "runner", runner)
					finished <- struct{}{}
					return
				}
			case <-ctx.Done():
				slog.Debug("gpu VRAM usage didn't recover within timeout", "seconds", time.Since(start).Seconds(), "free_before", format.HumanBytes2(freeMemoryBefore), "free_now", format.HumanBytes2(freeMemoryNow), "runner", runner)
719
720
721
722
723
724
725
726
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

727
728
729
730
731
732
733
734
735
736
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,
737
			slog.Any("inference", runner.gpus),
738
739
740
		)
	}
	attrs = append(attrs,
Jesse Gross's avatar
Jesse Gross committed
741
742
		slog.String("size", format.HumanBytes2(runner.totalSize)),
		slog.String("vram", format.HumanBytes2(runner.vramSize)),
743
744
745
746
		slog.Int("parallel", runner.numParallel),
		slog.Int("pid", runner.pid),
		slog.String("model", runner.modelPath),
	)
747
748
749
	if runner.Options != nil {
		attrs = append(attrs, slog.Int("num_ctx", runner.Options.NumCtx))
	}
750
751
752
	return slog.GroupValue(attrs...)
}

753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
// Implements discover.RunnerDiscovery
func (runner *runnerRef) GetPort() int {
	if runner.llama != nil {
		return runner.llama.GetPort()
	}
	return -1
}

func (runner *runnerRef) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo {
	if runner.llama != nil {
		return runner.llama.GetDeviceInfos(ctx)
	}
	return nil
}

func (runner *runnerRef) GetActiveDeviceIDs() []ml.DeviceID {
	return runner.gpus
}

func (runner *runnerRef) HasExited() bool {
	if runner.llama != nil {
		return runner.llama.HasExited()
	}
	return true
}

779
780
781
782
783
784
785
786
787
788
789
790
791
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
792
793
794
795
796
797
}

// 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] }
Jesse Gross's avatar
Jesse Gross committed
798
// func (a BySize) Less(i, j int) bool { return a[i].vramSize < a[j].vramSize }
799

Daniel Hiltgen's avatar
Daniel Hiltgen committed
800
// findRunnerToUnload finds a runner to unload to make room for a new model
801
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
802
803
804
805
806
807
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
808
809
810
811
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
812
813
814

	// 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?
815
	sort.Sort(ByDurationAndName(runnerList))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
816
817
818
819
820
821
822

	// 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 {
823
			slog.Debug("found an idle runner to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
824
825
826
827
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
828
	slog.Debug("no idle runners, picking the shortest duration", "runner_count", len(runnerList), "runner", runnerList[0])
Daniel Hiltgen's avatar
Daniel Hiltgen committed
829
830
831
832
833
834
	return runnerList[0]
}

func (s *Scheduler) unloadAllRunners() {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
Jesse Gross's avatar
Jesse Gross committed
835
836
837
838
839
840
841

	if s.activeLoading != nil {
		slog.Debug("shutting down currently loading runner")
		s.activeLoading.Close()
		s.activeLoading = nil
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
842
843
844
845
846
847
848
	for model, runner := range s.loaded {
		if runner.llama != nil {
			slog.Debug("shutting down runner", "model", model)
			runner.llama.Close()
		}
	}
}
849

Patrick Devine's avatar
Patrick Devine committed
850
851
852
func (s *Scheduler) expireRunner(model *Model) {
	s.loadedMu.Lock()
	runner, ok := s.loaded[model.ModelPath]
853
	s.loadedMu.Unlock()
Patrick Devine's avatar
Patrick Devine committed
854
855
856
857
858
859
860
861
862
863
864
865
866
867
	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()
	}
}