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

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

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

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

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

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

Jesse Gross's avatar
Jesse Gross committed
46
47
48
49
50
51
52
53
	// 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

	loadFn       func(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, requireFull bool) bool
Michael Yang's avatar
Michael Yang committed
54
	newServerFn  func(gpus discover.GpuInfoList, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int) (llm.LlamaServer, error)
55
56
	getGpuFn     func() discover.GpuInfoList
	getCpuFn     func() discover.GpuInfoList
57
	reschedDelay time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
58
59
}

60
61
62
63
64
// 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
65
var ErrMaxQueue = errors.New("server busy, please try again.  maximum pending requests exceeded")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
66
67

func InitScheduler(ctx context.Context) *Scheduler {
Michael Yang's avatar
int  
Michael Yang committed
68
	maxQueue := envconfig.MaxQueue()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
69
	sched := &Scheduler{
Michael Yang's avatar
int  
Michael Yang committed
70
71
72
		pendingReqCh:  make(chan *LlmRequest, maxQueue),
		finishedReqCh: make(chan *LlmRequest, maxQueue),
		expiredCh:     make(chan *runnerRef, maxQueue),
73
		unloadedCh:    make(chan any, maxQueue),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
74
75
		loaded:        make(map[string]*runnerRef),
		newServerFn:   llm.NewLlamaServer,
76
77
		getGpuFn:      discover.GetGPUInfo,
		getCpuFn:      discover.GetCPUInfo,
78
		reschedDelay:  250 * time.Millisecond,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
79
80
81
82
83
84
	}
	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
85
func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, sessionDuration *api.Duration) (chan *runnerRef, chan error) {
86
87
88
89
	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

Jesse Gross's avatar
Jesse Gross committed
90
91
92
93
94
	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
95
96
	req := &LlmRequest{
		ctx:             c,
Jesse Gross's avatar
Jesse Gross committed
97
		model:           m,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
98
99
		opts:            opts,
		sessionDuration: sessionDuration,
Jesse Gross's avatar
Jesse Gross committed
100
		successCh:       make(chan *runnerRef, 1),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
101
102
		errCh:           make(chan error, 1),
	}
103

Jesse Gross's avatar
Jesse Gross committed
104
105
106
107
108
109
110
111
112
113
114
	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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
	}
	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
139
			pending.schedAttempts++
140
141
142
143
144
145

			if pending.ctx.Err() != nil {
				slog.Debug("pending request cancelled or timed out, skipping scheduling")
				continue
			}

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

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

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

Jesse Gross's avatar
Jesse Gross committed
203
204
					// Update free memory from currently loaded models
					s.updateFreeSpace(gpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
205

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

Jesse Gross's avatar
Jesse Gross committed
213
214
					// More than one loaded model, so we have to see if the
					// new one fits
215

Jesse Gross's avatar
Jesse Gross committed
216
217
218
219
					needEvict := s.loadFn(pending, ggml, gpus, true)
					if !needEvict {
						slog.Debug("new model fits with existing models, loading")
						break
Daniel Hiltgen's avatar
Daniel Hiltgen committed
220
					}
Jesse Gross's avatar
Jesse Gross committed
221
222

					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
223
224
225
226
227
228
229
230
231
				}

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

321
			s.loadedMu.Lock()
322
323
324
325
326
327
328
329
330
331
332
			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)
333
334
335
336
337
338
339
340
341
342
343
			} 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()
344
345
346
347
348
349
350
351
352
353
354
355
			} else {
				slog.Debug("starting background wait for VRAM recovery", "runner", runner)
				finished := runner.waitForVRAMRecovery()
				runner.unload()
				delete(s.loaded, runner.modelPath)
				s.loadedMu.Unlock()
				slog.Debug("runner terminated and removed from list, blocking for VRAM recovery", "runner", runner)
				<-finished
				runner.refMu.Unlock()
				slog.Debug("sending an unloaded event", "runner", runner)
				s.unloadedCh <- struct{}{}
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
356
357
358
359
360
361
362
363
364
365
366
		}
	}
}

// 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++
367
368
369
370
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
371
372
373
	if pending.sessionDuration != nil {
		runner.sessionDuration = pending.sessionDuration.Duration
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
374
375
376
	pending.successCh <- runner
	go func() {
		<-pending.ctx.Done()
377
		slog.Debug("context for request finished", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
378
379
380
381
		finished <- pending
	}()
}

Jesse Gross's avatar
Jesse Gross committed
382
383
384
// 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.
func (s *Scheduler) load(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, requireFull bool) bool {
385
	numParallel := max(int(envconfig.NumParallel()), 1)
Jesse Gross's avatar
Jesse Gross committed
386
387
388
389
390
391
392
393
394
395
396
397
398

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

	// `mllama` is a snowflake and uses an encoder cache which cannot be used with num_parallel > 1
	// ref: https://github.com/ollama/ollama/issues/4165
	if slices.Contains(req.model.Config.ModelFamilies, "mllama") && numParallel != 1 {
		numParallel = 1
		slog.Warn("mllama does not currently support parallel requests")
	}

Michael Yang's avatar
Michael Yang committed
399
	sessionDuration := envconfig.KeepAlive()
400
401
402
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Jesse Gross's avatar
Jesse Gross committed
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432

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

	if llama == nil {
		var err error
		llama, err = s.newServerFn(gpus, req.model.ModelPath, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, numParallel)
		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()

	err := llama.Load(req.ctx, gpus, requireFull)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
433
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
434
435
		if errors.Is(err, llm.ErrLoadRequiredFull) {
			return true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
436
		}
Jesse Gross's avatar
Jesse Gross committed
437
438
439
440

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

445
446
447
448
449
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
450
		sessionDuration: sessionDuration,
451
		gpus:            gpus,
Jesse Gross's avatar
Jesse Gross committed
452
453
		vramSize:        llama.VRAMSize(),
		totalSize:       llama.TotalSize(),
454
		loading:         true,
455
		pid:             llama.Pid(),
456
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
457
	runner.numParallel = numParallel
458
	runner.refMu.Lock() // hold lock until running or aborted
459

Daniel Hiltgen's avatar
Daniel Hiltgen committed
460
	s.loadedMu.Lock()
461
462
463
464
465
466
467
	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
468
	s.activeLoading = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
469
470
471
472
473
474
475
476
477
	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
478
			slog.Debug("triggering expiration for failed load", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
479
480
481
			s.expiredCh <- runner
			return
		}
482
		slog.Debug("finished setting up", "runner", runner)
483
484
485
		if runner.pid < 0 {
			runner.pid = llama.Pid()
		}
486
		runner.refCount++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
487
488
489
490
491
492
493
494
		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
495
496

	return false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
497
498
}

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

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

Jesse Gross's avatar
Jesse Gross committed
547
548
549
550
551
552
	llama     llm.LlamaServer
	pid       int
	loading   bool                 // True only during initial load, then false forever
	gpus      discover.GpuInfoList // Recorded at time of provisioning
	vramSize  uint64
	totalSize uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
553
554
555

	sessionDuration time.Duration
	expireTimer     *time.Timer
556
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
557

Daniel Hiltgen's avatar
Daniel Hiltgen committed
558
559
560
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
561
562
563
564
565
	*api.Options
}

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

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

589
590
591
592
	if runner.Options == nil {
		return true
	}

593
594
595
596
597
598
599
600
601
	// 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
602
	defer cancel()
603
604
	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
605
606
607
608
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
609

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

613
614
615
616
617
618
619
620
// 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
621
622
func (runner *runnerRef) waitForVRAMRecovery() chan any {
	finished := make(chan any, 1)
623

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

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

			// Query GPUs, look for free to go back up
654
			gpusNow := discover.GetGPUInfo()
655
656
657
658
659
660
			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
Jesse Gross's avatar
Jesse Gross committed
661
			if float32(freeMemoryNow-freeMemoryBefore) > float32(runner.vramSize)*0.8 {
662
				slog.Debug(fmt.Sprintf("gpu VRAM free memory converged after %0.2f seconds", time.Since(start).Seconds()), "runner", runner)
663
664
665
666
667
668
669
670
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
func (runner *runnerRef) LogValue() slog.Value {
	if runner == nil {
		return slog.StringValue("nil")
	}
	attrs := []slog.Attr{}
	if runner.model != nil {
		attrs = append(attrs, slog.String("name", runner.model.Name))
	}
	if len(runner.gpus) > 0 {
		attrs = append(attrs,
			slog.String("inference", runner.gpus[0].Library),
			slog.Int("devices", len(runner.gpus)),
		)
	}
	attrs = append(attrs,
Jesse Gross's avatar
Jesse Gross committed
686
687
		slog.String("size", format.HumanBytes2(runner.totalSize)),
		slog.String("vram", format.HumanBytes2(runner.vramSize)),
688
689
690
691
		slog.Int("parallel", runner.numParallel),
		slog.Int("pid", runner.pid),
		slog.String("model", runner.modelPath),
	)
692
693
694
	if runner.Options != nil {
		attrs = append(attrs, slog.Int("num_ctx", runner.Options.NumCtx))
	}
695
696
697
	return slog.GroupValue(attrs...)
}

698
699
700
701
702
703
704
705
706
707
708
709
710
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
711
712
713
714
715
716
}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
719
// findRunnerToUnload finds a runner to unload to make room for a new model
720
func (s *Scheduler) findRunnerToUnload() *runnerRef {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
721
722
723
724
725
726
	s.loadedMu.Lock()
	runnerList := make([]*runnerRef, 0, len(s.loaded))
	for _, r := range s.loaded {
		runnerList = append(runnerList, r)
	}
	s.loadedMu.Unlock()
727
728
729
730
	if len(runnerList) == 0 {
		slog.Debug("no loaded runner to unload")
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
731
732
733

	// 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?
734
	sort.Sort(ByDurationAndName(runnerList))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
735
736
737
738
739
740
741

	// 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 {
742
			slog.Debug("found an idle runner to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
743
744
745
746
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
747
	slog.Debug("no idle runners, picking the shortest duration", "runner_count", len(runnerList), "runner", runnerList[0])
Daniel Hiltgen's avatar
Daniel Hiltgen committed
748
749
750
751
752
753
	return runnerList[0]
}

func (s *Scheduler) unloadAllRunners() {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
Jesse Gross's avatar
Jesse Gross committed
754
755
756
757
758
759
760

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
761
762
763
764
765
766
767
	for model, runner := range s.loaded {
		if runner.llama != nil {
			slog.Debug("shutting down runner", "model", model)
			runner.llama.Close()
		}
	}
}
768

Patrick Devine's avatar
Patrick Devine committed
769
770
771
func (s *Scheduler) expireRunner(model *Model) {
	s.loadedMu.Lock()
	runner, ok := s.loaded[model.ModelPath]
772
	s.loadedMu.Unlock()
Patrick Devine's avatar
Patrick Devine committed
773
774
775
776
777
778
779
780
781
782
783
784
785
786
	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()
	}
}