sched.go 28.3 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
	gpuIDs, err := llama.Load(req.ctx, systemInfo, gpus, requireFull)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
441
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
442
		if errors.Is(err, llm.ErrLoadRequiredFull) {
443
444
445
446
447
448
449
			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
450
			return true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
451
		}
Jesse Gross's avatar
Jesse Gross committed
452
453
454
455

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

460
461
462
463
464
465
466
467
468
469
470
471
472
473
	// 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
				}
			}
		}
	}

474
475
476
477
478
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
479
		sessionDuration: sessionDuration,
480
		gpus:            gpuIDs,
481
		discreteGPUs:    discreteGPUs,
Jesse Gross's avatar
Jesse Gross committed
482
483
		vramSize:        llama.VRAMSize(),
		totalSize:       llama.TotalSize(),
484
		loading:         true,
485
		pid:             llama.Pid(),
486
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
487
	runner.numParallel = numParallel
488
	runner.refMu.Lock() // hold lock until running or aborted
489

Daniel Hiltgen's avatar
Daniel Hiltgen committed
490
	s.loadedMu.Lock()
491
492
493
494
495
496
497
	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
498
	s.activeLoading = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
499
500
501
502
503
504
505
506
507
	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
508
			slog.Debug("triggering expiration for failed load", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
509
510
511
			s.expiredCh <- runner
			return
		}
512
		slog.Debug("finished setting up", "runner", runner)
513
514
515
		if runner.pid < 0 {
			runner.pid = llama.Pid()
		}
516
		runner.refCount++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
517
518
519
520
521
522
523
524
		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
525
526

	return false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
527
528
}

529
530
531
532
func (s *Scheduler) updateFreeSpace(allGpus []ml.DeviceInfo) {
	if len(allGpus) == 0 {
		return
	}
533
	predMap := map[ml.DeviceID]uint64{} // Sum up the total predicted usage per GPU for all runners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
534
	s.loadedMu.Lock()
535
	runners := make([]*runnerRef, 0, len(s.loaded))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
536
	for _, r := range s.loaded {
537
538
539
540
		runners = append(runners, r)
	}
	s.loadedMu.Unlock()
	for _, r := range runners {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
541
542
543
		r.refMu.Lock()
		if r.llama != nil {
			for _, gpu := range allGpus {
544
				predMap[gpu.DeviceID] += r.llama.VRAMByGPU(gpu.DeviceID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
545
546
547
548
549
550
551
552
553
			}
		} 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 {
554
		if p, ok := predMap[allGpus[i].DeviceID]; ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
555
556
557
558
559
560
561
562
			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
563
				// 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
564
565
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
566
567
568
569
570
571
			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
572
type runnerRef struct {
573
	refMu    sync.Mutex
Daniel Hiltgen's avatar
Daniel Hiltgen committed
574
575
	refCount uint // prevent unloading if > 0

576
577
578
579
580
581
582
	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
583
584
585

	sessionDuration time.Duration
	expireTimer     *time.Timer
586
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
587

Daniel Hiltgen's avatar
Daniel Hiltgen committed
588
589
590
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
591
592
593
594
595
	*api.Options
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
613
614
615
616
	timeout := 10 * time.Second
	if runner.loading {
		timeout = 2 * time.Minute // Initial load can take a long time for big models on slow systems...
	}
617

618
619
620
621
	if runner.Options == nil {
		return true
	}

622
623
624
625
626
627
628
629
630
	// 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
631
	defer cancel()
632
633
	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
634
635
636
637
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
638

Daniel Hiltgen's avatar
Daniel Hiltgen committed
639
640
641
	return false
}

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
650
func (s *Scheduler) waitForVRAMRecovery(runner *runnerRef, runners []ml.FilteredRunnerDiscovery) chan any {
651
	finished := make(chan any, 1)
652

653
654
655
	// 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") {
656
		finished <- struct{}{}
657
		slog.Debug("no need to wait for VRAM recovery", "runner", runner)
658
659
660
661
662
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
663
	gpusBefore := s.getGpuFn(context.Background(), runners)
664
665
666
667
668
	var totalMemoryBefore, freeMemoryBefore uint64
	for _, gpu := range gpusBefore {
		totalMemoryBefore += gpu.TotalMemory
		freeMemoryBefore += gpu.FreeMemory
	}
669
670
671
	totalMemoryNow := totalMemoryBefore
	freeMemoryNow := freeMemoryBefore

672
	go func() {
673
674
		// 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)
675
		defer cancel()
676
677
678
		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
679
680
681
682
683
684
685
686
687
688
			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
				}
689
690
691
692
693
				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)
				}
694
695
696
697
698
699
700
701
				// 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)
702
703
704
705
706
707
708
709
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

710
711
712
713
714
715
716
717
718
719
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,
720
			slog.Any("inference", runner.gpus),
721
722
723
		)
	}
	attrs = append(attrs,
Jesse Gross's avatar
Jesse Gross committed
724
725
		slog.String("size", format.HumanBytes2(runner.totalSize)),
		slog.String("vram", format.HumanBytes2(runner.vramSize)),
726
727
728
729
		slog.Int("parallel", runner.numParallel),
		slog.Int("pid", runner.pid),
		slog.String("model", runner.modelPath),
	)
730
731
732
	if runner.Options != nil {
		attrs = append(attrs, slog.Int("num_ctx", runner.Options.NumCtx))
	}
733
734
735
	return slog.GroupValue(attrs...)
}

736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
// 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
}

762
763
764
765
766
767
768
769
770
771
772
773
774
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
775
776
777
778
779
780
}

// 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
781
// func (a BySize) Less(i, j int) bool { return a[i].vramSize < a[j].vramSize }
782

Daniel Hiltgen's avatar
Daniel Hiltgen committed
783
// findRunnerToUnload finds a runner to unload to make room for a new model
784
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
785
786
787
788
789
790
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
791
792
793
794
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
795
796
797

	// 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?
798
	sort.Sort(ByDurationAndName(runnerList))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
799
800
801
802
803
804
805

	// 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 {
806
			slog.Debug("found an idle runner to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
807
808
809
810
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
811
	slog.Debug("no idle runners, picking the shortest duration", "runner_count", len(runnerList), "runner", runnerList[0])
Daniel Hiltgen's avatar
Daniel Hiltgen committed
812
813
814
815
816
817
	return runnerList[0]
}

func (s *Scheduler) unloadAllRunners() {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
Jesse Gross's avatar
Jesse Gross committed
818
819
820
821
822
823
824

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
825
826
827
828
829
830
831
	for model, runner := range s.loaded {
		if runner.llama != nil {
			slog.Debug("shutting down runner", "model", model)
			runner.llama.Close()
		}
	}
}
832

Patrick Devine's avatar
Patrick Devine committed
833
834
835
func (s *Scheduler) expireRunner(model *Model) {
	s.loadedMu.Lock()
	runner, ok := s.loaded[model.ModelPath]
836
	s.loadedMu.Unlock()
Patrick Devine's avatar
Patrick Devine committed
837
838
839
840
841
842
843
844
845
846
847
848
849
850
	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()
	}
}