sched.go 27.2 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/logutil"
25
	"github.com/ollama/ollama/ml"
26
	"github.com/ollama/ollama/types/model"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
27
28
29
30
31
32
)

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

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

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

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

55
56
57
58
59
60
61
	loadFn      func(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList, requireFull bool) bool
	newServerFn func(gpus discover.GpuInfoList, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int) (llm.LlamaServer, error)
	getGpuFn    func(ctx context.Context, runners []discover.FilteredRunnerDiscovery) discover.GpuInfoList
	getCpuFn    func() discover.GpuInfo

	// waitForRecovery sets the limit for how long to wait for memory usage to recover after unload before scheduling the next model
	waitForRecovery time.Duration
Daniel Hiltgen's avatar
Daniel Hiltgen committed
62
63
}

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

func InitScheduler(ctx context.Context) *Scheduler {
Michael Yang's avatar
int  
Michael Yang committed
72
	maxQueue := envconfig.MaxQueue()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
73
	sched := &Scheduler{
74
75
76
77
78
79
80
81
82
		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,
		getGpuFn:        discover.GetGPUInfo,
		getCpuFn:        discover.GetCPUInfo,
		waitForRecovery: 5 * time.Second,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
83
84
85
86
87
88
	}
	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
89
func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, sessionDuration *api.Duration) (chan *runnerRef, chan error) {
90
91
92
93
	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

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

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

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

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
161
162
				if runner != nil {
					if runner.needsReload(ctx, pending) {
163
						slog.Debug("reloading", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
164
165
166
167
168
169
						runnerToExpire = runner
					} else {
						// Runner is usable, return it
						pending.useLoadedRunner(runner, s.finishedReqCh)
						break
					}
Michael Yang's avatar
Michael Yang committed
170
				} else if envconfig.MaxRunners() > 0 && loadedCount >= int(envconfig.MaxRunners()) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
171
					slog.Debug("max runners achieved, unloading one to make room", "runner_count", loadedCount)
172
					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
173
				} else {
174
					// Either no models are loaded or below envconfig.MaxRunners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
175
					// Get a refreshed GPU list
176
					var gpus discover.GpuInfoList
177
					if pending.opts.NumGPU == 0 {
178
						gpus = discover.GpuInfoList{s.getCpuFn()}
179
					} else {
180
						gpus = s.getGpuFn(ctx, runnersSnapshot)
181
					}
182

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

205
					// Load model for fitting
206
					ggml, err := llm.LoadModel(pending.model.ModelPath, 1024)
207
208
209
210
					if err != nil {
						pending.errCh <- err
						break
					}
211

Jesse Gross's avatar
Jesse Gross committed
212
213
					// Update free memory from currently loaded models
					s.updateFreeSpace(gpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
214

Jesse Gross's avatar
Jesse Gross committed
215
					if loadedCount == 0 {
216
						// No models loaded. Load the model but prefer the best fit.
217
						slog.Debug("loading first model", "model", pending.model.ModelPath)
Jesse Gross's avatar
Jesse Gross committed
218
						s.loadFn(pending, ggml, gpus, false)
219
220
221
						break
					}

Jesse Gross's avatar
Jesse Gross committed
222
223
					// More than one loaded model, so we have to see if the
					// new one fits
224

Jesse Gross's avatar
Jesse Gross committed
225
226
227
228
					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
229
					}
Jesse Gross's avatar
Jesse Gross committed
230
231

					runnerToExpire = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
232
233
234
				}

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

331
			s.loadedMu.Lock()
332
333
334
335
336
337
338
339
340
341
342
			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)
343
344
345
346
347
348
349
350
351
352
353
			} 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()
354
355
			} else {
				slog.Debug("starting background wait for VRAM recovery", "runner", runner)
356
357
358
359
360
				runnersSnapshot := make([]discover.FilteredRunnerDiscovery, 0, len(s.loaded))
				for _, r := range s.loaded {
					runnersSnapshot = append(runnersSnapshot, r)
				}
				finished := s.waitForVRAMRecovery(runner, runnersSnapshot)
361
362
363
364
365
366
367
368
369
				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
370
371
372
373
374
375
376
377
378
379
380
		}
	}
}

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

Jesse Gross's avatar
Jesse Gross committed
396
397
398
// 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 {
399
	numParallel := max(int(envconfig.NumParallel()), 1)
Jesse Gross's avatar
Jesse Gross committed
400
401
402
403
404
405
406
407
408
409
410
411
412

	// 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
413
	sessionDuration := envconfig.KeepAlive()
414
415
416
	if req.sessionDuration != nil {
		sessionDuration = req.sessionDuration.Duration
	}
Jesse Gross's avatar
Jesse Gross committed
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445

	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()

446
	gpuIDs, err := llama.Load(req.ctx, gpus, requireFull)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
447
	if err != nil {
Jesse Gross's avatar
Jesse Gross committed
448
449
		if errors.Is(err, llm.ErrLoadRequiredFull) {
			return true
Daniel Hiltgen's avatar
Daniel Hiltgen committed
450
		}
Jesse Gross's avatar
Jesse Gross committed
451
452
453
454

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

459
460
461
462
463
	runner := &runnerRef{
		model:           req.model,
		modelPath:       req.model.ModelPath,
		llama:           llama,
		Options:         &req.opts,
464
		sessionDuration: sessionDuration,
465
		gpus:            gpuIDs,
Jesse Gross's avatar
Jesse Gross committed
466
467
		vramSize:        llama.VRAMSize(),
		totalSize:       llama.TotalSize(),
468
		loading:         true,
469
		pid:             llama.Pid(),
470
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
471
	runner.numParallel = numParallel
472
	runner.refMu.Lock() // hold lock until running or aborted
473

Daniel Hiltgen's avatar
Daniel Hiltgen committed
474
	s.loadedMu.Lock()
475
476
477
478
479
480
481
	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
482
	s.activeLoading = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
483
484
485
486
487
488
489
490
491
	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
492
			slog.Debug("triggering expiration for failed load", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
493
494
495
			s.expiredCh <- runner
			return
		}
496
		slog.Debug("finished setting up", "runner", runner)
497
498
499
		if runner.pid < 0 {
			runner.pid = llama.Pid()
		}
500
		runner.refCount++
Daniel Hiltgen's avatar
Daniel Hiltgen committed
501
502
503
504
505
506
507
508
		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
509
510

	return false
Daniel Hiltgen's avatar
Daniel Hiltgen committed
511
512
}

513
func (s *Scheduler) updateFreeSpace(allGpus discover.GpuInfoList) {
514
	predMap := map[ml.DeviceID]uint64{} // Sum up the total predicted usage per GPU for all runners
Daniel Hiltgen's avatar
Daniel Hiltgen committed
515
	s.loadedMu.Lock()
516
	runners := make([]*runnerRef, 0, len(s.loaded))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
517
	for _, r := range s.loaded {
518
519
520
521
		runners = append(runners, r)
	}
	s.loadedMu.Unlock()
	for _, r := range runners {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
522
523
524
		r.refMu.Lock()
		if r.llama != nil {
			for _, gpu := range allGpus {
525
				predMap[gpu.DeviceID] += r.llama.VRAMByGPU(gpu.DeviceID)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
526
527
528
529
530
531
532
533
534
			}
		} 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 {
535
		if p, ok := predMap[allGpus[i].DeviceID]; ok {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
536
537
538
539
540
541
542
543
			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
544
				// 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
545
546
				allGpus[i].FreeMemory = allGpus[i].TotalMemory - p
			}
547
548
549
550
551
552
			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
553
type runnerRef struct {
554
	refMu    sync.Mutex
Daniel Hiltgen's avatar
Daniel Hiltgen committed
555
556
	refCount uint // prevent unloading if > 0

Jesse Gross's avatar
Jesse Gross committed
557
558
	llama     llm.LlamaServer
	pid       int
559
560
	loading   bool          // True only during initial load, then false forever
	gpus      []ml.DeviceID // Recorded at time of provisioning
Jesse Gross's avatar
Jesse Gross committed
561
562
	vramSize  uint64
	totalSize uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
563
564
565

	sessionDuration time.Duration
	expireTimer     *time.Timer
566
	expiresAt       time.Time
Daniel Hiltgen's avatar
Daniel Hiltgen committed
567

Daniel Hiltgen's avatar
Daniel Hiltgen committed
568
569
570
	model       *Model
	modelPath   string
	numParallel int
Daniel Hiltgen's avatar
Daniel Hiltgen committed
571
572
573
574
575
	*api.Options
}

// The refMu must already be held when calling unload
func (runner *runnerRef) unload() {
576
577
578
579
	if runner.expireTimer != nil {
		runner.expireTimer.Stop()
		runner.expireTimer = nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
580
581
582
	if runner.llama != nil {
		runner.llama.Close()
	}
583
	runner.model = nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
584
585
586
587
588
589
590
591
	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()
592

Daniel Hiltgen's avatar
Daniel Hiltgen committed
593
594
595
596
	timeout := 10 * time.Second
	if runner.loading {
		timeout = 2 * time.Minute // Initial load can take a long time for big models on slow systems...
	}
597

598
599
600
601
	if runner.Options == nil {
		return true
	}

602
603
604
605
606
607
608
609
610
	// 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
611
	defer cancel()
612
613
	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
614
615
616
617
		!reflect.DeepEqual(optsExisting, optsNew) || // have the runner options changed?
		runner.llama.Ping(ctx) != nil {
		return true
	}
618

Daniel Hiltgen's avatar
Daniel Hiltgen committed
619
620
621
	return false
}

622
623
624
625
626
627
628
629
// 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
630
func (s *Scheduler) waitForVRAMRecovery(runner *runnerRef, runners []discover.FilteredRunnerDiscovery) chan any {
631
	finished := make(chan any, 1)
632

633
634
	// CPU or Metal don't need checking, so no waiting required
	// windows can page VRAM, only cuda currently can report accurate used vram usage
635
	if len(runner.gpus) == 0 ||
636
637
		(len(runner.gpus) == 1 && (runner.gpus[0].Library == "cpu" || runner.gpus[0].Library == "Metal")) ||
		(runtime.GOOS == "windows" && runner.gpus[0].Library != "CUDA") {
638
		finished <- struct{}{}
639
		slog.Debug("no need to wait for VRAM recovery", "runner", runner)
640
641
642
643
644
		return finished
	}
	start := time.Now()

	// Establish a baseline before we unload
645
	gpusBefore := s.getGpuFn(context.Background(), runners)
646
647
648
649
650
	var totalMemoryBefore, freeMemoryBefore uint64
	for _, gpu := range gpusBefore {
		totalMemoryBefore += gpu.TotalMemory
		freeMemoryBefore += gpu.FreeMemory
	}
651
652
653
	totalMemoryNow := totalMemoryBefore
	freeMemoryNow := freeMemoryBefore

654
	go func() {
655
656
		// 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)
657
		defer cancel()
658
659
660
		ticker := time.NewTicker(250 * time.Millisecond)
		defer ticker.Stop()
		for {
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
			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
				}
				logutil.Trace("gpu VRAM convergence", "percent", int(max(float32(freeMemoryNow-freeMemoryBefore), 0.0)/float32(runner.vramSize)*100))
				// 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)
680
681
682
683
684
685
686
687
				finished <- struct{}{}
				return
			}
		}
	}()
	return finished
}

688
689
690
691
692
693
694
695
696
697
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,
698
			slog.Any("inference", runner.gpus),
699
700
701
		)
	}
	attrs = append(attrs,
Jesse Gross's avatar
Jesse Gross committed
702
703
		slog.String("size", format.HumanBytes2(runner.totalSize)),
		slog.String("vram", format.HumanBytes2(runner.vramSize)),
704
705
706
707
		slog.Int("parallel", runner.numParallel),
		slog.Int("pid", runner.pid),
		slog.String("model", runner.modelPath),
	)
708
709
710
	if runner.Options != nil {
		attrs = append(attrs, slog.Int("num_ctx", runner.Options.NumCtx))
	}
711
712
713
	return slog.GroupValue(attrs...)
}

714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
// 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
}

740
741
742
743
744
745
746
747
748
749
750
751
752
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
753
754
755
756
757
758
}

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

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

	// 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?
776
	sort.Sort(ByDurationAndName(runnerList))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
777
778
779
780
781
782
783

	// 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 {
784
			slog.Debug("found an idle runner to unload", "runner", runner)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
785
786
787
788
			return runner
		}
	}
	// None appear idle, just wait for the one with the shortest duration
789
	slog.Debug("no idle runners, picking the shortest duration", "runner_count", len(runnerList), "runner", runnerList[0])
Daniel Hiltgen's avatar
Daniel Hiltgen committed
790
791
792
793
794
795
	return runnerList[0]
}

func (s *Scheduler) unloadAllRunners() {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
Jesse Gross's avatar
Jesse Gross committed
796
797
798
799
800
801
802

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
803
804
805
806
807
808
809
	for model, runner := range s.loaded {
		if runner.llama != nil {
			slog.Debug("shutting down runner", "model", model)
			runner.llama.Close()
		}
	}
}
810

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