baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

unload in critical section (#4187)

Browse source
This commit is contained in:
Jeffrey Morgan 2024-05-05 17:18:27 -07:00 committed by GitHub
parent 840424a2c4
commit dfa2f32ca0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
server/sched.go
View file

@ -116,7 +116,7 @@ func (s *Scheduler) processPending(ctx context.Context) {
}
} else if envconfig.MaxRunners > 0 && loadedCount >= envconfig.MaxRunners {
slog.Debug("max runners achieved, unloading one to make room", "runner_count", loadedCount)
runnerToExpire = s.findRunnerToUnload(pending)
runnerToExpire = s.findRunnerToUnload()
} else {
// Either no models are loaded or below envconfig.MaxRunners
// Get a refreshed GPU list
@ -157,7 +157,7 @@ func (s *Scheduler) processPending(ctx context.Context) {
s.loadFn(pending, ggml, gpus)
break
}
runnerToExpire = s.findRunnerToUnload(pending)
runnerToExpire = s.findRunnerToUnload()
}
if runnerToExpire == nil {
@ -257,9 +257,9 @@ func (s *Scheduler) processCompleted(ctx context.Context) {
continue
}
s.loadedMu.Lock()
slog.Debug("got lock to unload", "model", runner.model)
runner.unload()
s.loadedMu.Lock()
delete(s.loaded, runner.model)
s.loadedMu.Unlock()
slog.Debug("runner released", "model", runner.model)
@ -504,7 +504,7 @@ func pickBestFitGPUs(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList) gpu.
}
// findRunnerToUnload finds a runner to unload to make room for a new model
func (s *Scheduler) findRunnerToUnload(req *LlmRequest) *runnerRef {
func (s *Scheduler) findRunnerToUnload() *runnerRef {
s.loadedMu.Lock()
runnerList := make([]*runnerRef, 0, len(s.loaded))
for _, r := range s.loaded {

9
server/sched_test.go
View file

@ -473,10 +473,7 @@ func TestUpdateFreeSpace(t *testing.T) {
func TestFindRunnerToUnload(t *testing.T) {
ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer done()
req := &LlmRequest{
ctx: ctx,
opts: api.DefaultOptions(),
}
r1 := &runnerRef{refCount: 1, sessionDuration: 1}
r2 := &runnerRef{sessionDuration: 2}
@ -486,10 +483,10 @@ func TestFindRunnerToUnload(t *testing.T) {
s.loaded["b"] = r2
s.loadedMu.Unlock()
resp := s.findRunnerToUnload(req)
resp := s.findRunnerToUnload()
require.Equal(t, r2, resp)
r2.refCount = 1
resp = s.findRunnerToUnload(req)
resp = s.findRunnerToUnload()
require.Equal(t, r1, resp)
}