package server import ( "bytes" "context" "errors" "log/slog" "os" "sync" "testing" "time" "github.com/stretchr/testify/require" "github.com/ollama/ollama/api" "github.com/ollama/ollama/format" "github.com/ollama/ollama/fs/ggml" "github.com/ollama/ollama/llm" "github.com/ollama/ollama/ml" "github.com/ollama/ollama/types/model" ) func TestMain(m *testing.M) { os.Setenv("OLLAMA_DEBUG", "1") logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})) slog.SetDefault(logger) os.Exit(m.Run()) } func TestSchedInit(t *testing.T) { ctx, done := context.WithCancel(t.Context()) defer done() s := InitScheduler(ctx) s.loadedMu.Lock() require.NotNil(t, s.loaded) s.loadedMu.Unlock() } func TestSchedLoad(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 20*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond modelPath, _ := createBinFile(t, ggml.KV{ "general.architecture": "llama", "llama.context_length": uint32(32), "llama.embedding_length": uint32(4096), "llama.block_count": uint32(1), "llama.attention.head_count": uint32(32), "llama.attention.head_count_kv": uint32(32), "tokenizer.ggml.tokens": []string{" "}, "tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.token_type": []int32{0}, }, []*ggml.Tensor{ {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, }) req := &LlmRequest{ ctx: ctx, model: &Model{ModelPath: modelPath}, opts: api.DefaultOptions(), successCh: make(chan *runnerRef, 1), errCh: make(chan error, 1), sessionDuration: &api.Duration{Duration: 2 * time.Second}, } // Fail to load model first s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return nil, errors.New("something failed to load model blah") } gpus := []ml.DeviceInfo{} systemInfo := ml.SystemInfo{} s.load(req, systemInfo, gpus, false) require.Empty(t, req.successCh) require.Len(t, req.errCh, 1) s.loadedMu.Lock() require.Empty(t, s.loaded) s.loadedMu.Unlock() err := <-req.errCh require.Contains(t, err.Error(), "this model may be incompatible") server := &mockLlm{vramSize: 10, vramByGPU: map[ml.DeviceID]uint64{}} s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { server.modelPath = model return server, nil } s.load(req, systemInfo, gpus, false) select { case err := <-req.errCh: require.NoError(t, err) case resp := <-req.successCh: require.Equal(t, uint64(10), resp.vramSize) require.Equal(t, uint(1), resp.refCount) s.loadedMu.Lock() require.Len(t, s.loaded, 1) s.loadedMu.Unlock() } modelPath2, _ := createBinFile(t, ggml.KV{ "general.architecture": "llama", "llama.context_length": uint32(32), "llama.embedding_length": uint32(4096), "llama.block_count": uint32(1), "llama.attention.head_count": uint32(32), "llama.attention.head_count_kv": uint32(32), "tokenizer.ggml.tokens": []string{" "}, "tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.token_type": []int32{0}, }, []*ggml.Tensor{ {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, }) req.model.ModelPath = modelPath2 server.waitResp = errors.New("wait failure") s.load(req, systemInfo, gpus, false) select { case err := <-req.errCh: require.Contains(t, err.Error(), "wait failure") case resp := <-req.successCh: t.Fatalf("unexpected success %v", resp) } s.loadedMu.Lock() runner := s.loaded[modelPath2] s.loadedMu.Unlock() require.NotNil(t, runner) require.Equal(t, uint(0), runner.refCount) time.Sleep(1 * time.Millisecond) require.Len(t, s.expiredCh, 1) } func TestSchedLoadStoresEffectiveContextLength(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) scenario := newScenarioRequest(t, ctx, "test", 10, nil, map[ml.DeviceID]uint64{}) scenario.req.opts.NumCtx = 262144 scenario.req.numCtxAuto = true scenario.srv.contextLength = 131072 s.newServerFn = scenario.newServer s.load(scenario.req, ml.SystemInfo{}, nil, false) select { case err := <-scenario.req.errCh: require.NoError(t, err) case runner := <-scenario.req.successCh: require.Equal(t, 131072, runner.Options.NumCtx) } } func TestSchedLoadStoresEffectiveExplicitContextLength(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) scenario := newScenarioRequest(t, ctx, "test", 10, nil, map[ml.DeviceID]uint64{}) scenario.req.opts.NumCtx = 262144 scenario.srv.contextLength = 131072 s.newServerFn = scenario.newServer s.load(scenario.req, ml.SystemInfo{}, nil, false) select { case err := <-scenario.req.errCh: require.NoError(t, err) case runner := <-scenario.req.successCh: require.Equal(t, 131072, runner.Options.NumCtx) } } func TestSchedVisionContextFloor(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() visionModel := &Model{ Name: "vision-test", Config: model.ConfigV2{ Capabilities: []string{string(model.CapabilityVision)}, }, } t.Run("automatic num_ctx is floored", func(t *testing.T) { s := InitScheduler(ctx) opts := api.DefaultOptions() opts.NumCtx = 128 s.getRunner(ctx, visionModel, opts, nil, true, false, nil) req := <-s.pendingReqCh require.Equal(t, 2048, req.opts.NumCtx) require.True(t, req.numCtxAuto) }) t.Run("explicit num_ctx is floored", func(t *testing.T) { s := InitScheduler(ctx) opts := api.DefaultOptions() opts.NumCtx = 128 s.getRunner(ctx, visionModel, opts, nil, false, false, nil) req := <-s.pendingReqCh require.Equal(t, 2048, req.opts.NumCtx) require.False(t, req.numCtxAuto) }) } type reqBundle struct { ctx context.Context //nolint:containedctx ctxDone func() srv *mockLlm req *LlmRequest } func (scenario *reqBundle) newServer(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { scenario.srv.modelPath = model return scenario.srv, nil } func newScenarioRequest(t *testing.T, ctx context.Context, modelName string, vramSize uint64, duration *api.Duration, vramByGPU map[ml.DeviceID]uint64) *reqBundle { return newScenarioRequestWithContext(t, ctx, modelName, vramSize, duration, vramByGPU, 32) } func newScenarioRequestWithContext(t *testing.T, ctx context.Context, modelName string, vramSize uint64, duration *api.Duration, vramByGPU map[ml.DeviceID]uint64, trainCtx uint32) *reqBundle { b := &reqBundle{} b.ctx, b.ctxDone = context.WithCancel(ctx) t.Helper() p, _ := createBinFile(t, ggml.KV{ "general.architecture": "llama", "llama.context_length": trainCtx, "llama.embedding_length": uint32(4096), "llama.block_count": uint32(1), "llama.attention.head_count": uint32(32), "llama.attention.head_count_kv": uint32(32), "tokenizer.ggml.tokens": []string{" "}, "tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.token_type": []int32{0}, }, []*ggml.Tensor{ {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, }) model := &Model{Name: modelName, ModelPath: p} if duration == nil { duration = &api.Duration{Duration: 5 * time.Millisecond} } b.req = &LlmRequest{ ctx: b.ctx, model: model, opts: api.DefaultOptions(), sessionDuration: duration, successCh: make(chan *runnerRef, 1), errCh: make(chan error, 1), } b.srv = &mockLlm{vramSize: vramSize, vramByGPU: vramByGPU} return b } func getGpuFn(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { slog.Info("test getGpuFn called", "runners", runners) g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 12 * format.GigaByte return []ml.DeviceInfo{g} } func getSystemInfoFn() ml.SystemInfo { slog.Info("test getSystemInfoFn called") return ml.SystemInfo{ TotalMemory: 32 * format.GigaByte, FreeMemory: 26 * format.GigaByte, } } func TestSchedRequestsSameModelSameRequest(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = getGpuFn s.getSystemInfoFn = getSystemInfoFn a := newScenarioRequest(t, ctx, "ollama-model-1", 10, &api.Duration{Duration: 5 * time.Millisecond}, nil) b := newScenarioRequest(t, ctx, "ollama-model-1", 11, &api.Duration{Duration: 0}, nil) b.req.model = a.req.model s.newServerFn = a.newServer slog.Info("a") s.pendingReqCh <- a.req require.Len(t, s.pendingReqCh, 1) s.Run(ctx) select { case resp := <-a.req.successCh: require.Equal(t, resp.llama, a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, a.req.errCh) case err := <-a.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } // Same runner as first request due to not needing a reload s.newServerFn = b.newServer slog.Info("b") s.pendingReqCh <- b.req select { case resp := <-b.req.successCh: require.Equal(t, resp.llama, a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, b.req.errCh) case err := <-b.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } } func TestSchedRequestsSimpleReloadSameModel(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 5000*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 12 * format.GigaByte gMu := sync.Mutex{} s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { slog.Info("test getGpuFn called", "runners", runners) gMu.Lock() defer gMu.Unlock() return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn a := newScenarioRequest(t, ctx, "ollama-model-1", 10, &api.Duration{Duration: 5 * time.Millisecond}, nil) b := newScenarioRequest(t, ctx, "ollama-model-1", 20, &api.Duration{Duration: 5 * time.Millisecond}, nil) tmpModel := *a.req.model b.req.model = &tmpModel s.newServerFn = a.newServer slog.Info("a") s.pendingReqCh <- a.req require.Len(t, s.pendingReqCh, 1) s.Run(ctx) select { case resp := <-a.req.successCh: require.Equal(t, resp.llama, a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, a.req.errCh) case err := <-a.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } // Trigger a reload s.newServerFn = b.newServer b.req.model.AdapterPaths = []string{"new"} slog.Info("b") s.pendingReqCh <- b.req // finish first two requests, so model can reload time.Sleep(1 * time.Millisecond) a.ctxDone() // Report recovered VRAM usage time.Sleep(1 * time.Millisecond) gMu.Lock() g.FreeMemory = 24 * format.GigaByte gMu.Unlock() select { case resp := <-b.req.successCh: require.Equal(t, resp.llama, b.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, b.req.errCh) case err := <-b.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } } func TestSchedRequestsMultipleLoadedModels(t *testing.T) { slog.Info("TestRequestsMultipleLoadedModels") ctx, done := context.WithTimeout(t.Context(), 1000*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 12 * format.GigaByte gMu := sync.Mutex{} s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { slog.Info("test getGpuFn called", "runners", runners) gMu.Lock() defer gMu.Unlock() return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn // Multiple loaded models a := newScenarioRequest(t, ctx, "model-a-1g-gpu", 1*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 1 * format.GigaByte}) a.req.sessionDuration = &api.Duration{Duration: 5 * time.Millisecond} b := newScenarioRequest(t, ctx, "model-b-10g-gpu", 10*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 10 * format.GigaByte}) b.req.sessionDuration = &api.Duration{Duration: 5 * time.Millisecond} c := newScenarioRequest(t, ctx, "model-c-10g-cpu", 10*format.GigaByte, nil, nil /* No GPU load */) c.req.opts.NumGPU = 0 // CPU load, will be allowed b.req.sessionDuration = &api.Duration{Duration: 10 * time.Millisecond} // longer than b to cause the scheduler to favor unloading b over c d := newScenarioRequest(t, ctx, "model-d-10g-gpu", 13*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 13 * format.GigaByte}) // Needs prior unloaded s.newServerFn = a.newServer slog.Info("Loading A") s.pendingReqCh <- a.req s.Run(ctx) select { case resp := <-a.req.successCh: require.Equal(t, resp.llama, a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, a.req.errCh) case err := <-a.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } s.loadedMu.Lock() require.Len(t, s.loaded, 1) s.loadedMu.Unlock() t.Setenv("OLLAMA_MAX_LOADED_MODELS", "0") s.newServerFn = b.newServer slog.Info("Loading B") s.pendingReqCh <- b.req select { case resp := <-b.req.successCh: require.Equal(t, resp.llama, b.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, b.req.errCh) case err := <-b.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } s.loadedMu.Lock() require.Len(t, s.loaded, 2) s.loadedMu.Unlock() // This is a CPU load with NumGPU = 0 so it should load s.newServerFn = c.newServer slog.Info("Loading C") s.pendingReqCh <- c.req select { case resp := <-c.req.successCh: require.Equal(t, resp.llama, c.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, c.req.errCh) case err := <-c.req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): slog.Info("FAIL: scheduler state", "s.loaded", s.loaded) t.Fatal("timeout") } s.loadedMu.Lock() require.Len(t, s.loaded, 3) s.loadedMu.Unlock() // Try to load a model that won't fit s.newServerFn = d.newServer slog.Info("d") s.loadedMu.Lock() require.Len(t, s.loaded, 3) s.loadedMu.Unlock() a.ctxDone() // Won't help since this one isn't big enough to make room time.Sleep(2 * time.Millisecond) s.pendingReqCh <- d.req // finish prior request, so new model can load time.Sleep(6 * time.Millisecond) s.loadedMu.Lock() require.Len(t, s.loaded, 2) s.loadedMu.Unlock() // Mark b done so it can unload b.ctxDone() // Report recovered VRAM usage so scheduler will finish waiting and unload time.Sleep(1 * time.Millisecond) gMu.Lock() g.FreeMemory = 24 * format.GigaByte gMu.Unlock() select { case resp := <-d.req.successCh: require.Equal(t, resp.llama, d.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, d.req.errCh) case <-ctx.Done(): t.Fatal("timeout") } // Wait for b to close closeWait: for { select { case <-ctx.Done(): t.Fatal("timeout") default: if b.srv.closeCalled { break closeWait } time.Sleep(1 * time.Millisecond) } } s.loadedMu.Lock() require.Len(t, s.loaded, 2) s.loadedMu.Unlock() } func TestSchedGetRunner(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 3*time.Second) defer done() a := newScenarioRequest(t, ctx, "ollama-model-1a", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil) b := newScenarioRequest(t, ctx, "ollama-model-1b", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil) c := newScenarioRequest(t, ctx, "ollama-model-1c", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil) t.Setenv("OLLAMA_MAX_QUEUE", "1") s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = getGpuFn s.getSystemInfoFn = getSystemInfoFn s.newServerFn = a.newServer slog.Info("a") successCh1a, errCh1a := s.GetRunner(a.ctx, a.req.model, a.req.opts, a.req.sessionDuration) require.Len(t, s.pendingReqCh, 1) slog.Info("b") successCh1b, errCh1b := s.GetRunner(b.ctx, b.req.model, b.req.opts, b.req.sessionDuration) require.Len(t, s.pendingReqCh, 1) require.Empty(t, successCh1b) require.Len(t, errCh1b, 1) err := <-errCh1b require.Contains(t, err.Error(), "server busy") s.Run(ctx) select { case resp := <-successCh1a: require.Equal(t, resp.llama, a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, errCh1a) case err := <-errCh1a: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } a.ctxDone() // Set "a" model to idle so it can unload s.loadedMu.Lock() require.Len(t, s.loaded, 1) s.loadedMu.Unlock() c.req.model.ModelPath = "bad path" slog.Info("c") successCh1c, errCh1c := s.GetRunner(c.ctx, c.req.model, c.req.opts, c.req.sessionDuration) // Starts in pending channel, then should be quickly processed to return an error time.Sleep(50 * time.Millisecond) // Long enough for the "a" model to expire and unload require.Empty(t, successCh1c) s.loadedMu.Lock() require.Empty(t, s.loaded) s.loadedMu.Unlock() require.Len(t, errCh1c, 1) err = <-errCh1c require.Contains(t, err.Error(), "bad path") b.ctxDone() } func TestSchedGetRunnerUsesDigestKeyWhenModelPathEmpty(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() s := InitScheduler(ctx) opts := api.DefaultOptions() opts.NumCtx = 4 loadedModel := &Model{Name: "safetensors-a", Digest: "sha-a"} loadedRunner := &runnerRef{ model: loadedModel, modelKey: schedulerModelKey(loadedModel), llama: &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}, Options: &opts, numParallel: 1, } s.loadedMu.Lock() s.loaded[loadedRunner.modelKey] = loadedRunner s.loadedMu.Unlock() reqModel := &Model{Name: "safetensors-b", Digest: "sha-b"} successCh, errCh := s.GetRunner(ctx, reqModel, opts, nil) require.Empty(t, successCh) require.Empty(t, errCh) require.Len(t, s.pendingReqCh, 1) } func TestSchedGetRunnerReusesSameDigestWhenModelPathEmpty(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() s := InitScheduler(ctx) opts := api.DefaultOptions() opts.NumCtx = 4 loadedModel := &Model{Name: "safetensors-a", Digest: "sha-a"} loadedRunner := &runnerRef{ model: loadedModel, modelKey: schedulerModelKey(loadedModel), llama: &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}, Options: &opts, numParallel: 1, } s.loadedMu.Lock() s.loaded[loadedRunner.modelKey] = loadedRunner s.loadedMu.Unlock() reqCtx, cancelReq := context.WithCancel(ctx) successCh, errCh := s.GetRunner(reqCtx, &Model{Name: "safetensors-a-copy", Digest: "sha-a"}, opts, nil) cancelReq() select { case runner := <-successCh: require.Equal(t, loadedRunner, runner) default: t.Fatal("expected existing runner to be reused") } require.Empty(t, errCh) require.Empty(t, s.pendingReqCh) } func TestSchedExpireRunner(t *testing.T) { ctx, done := context.WithCancel(t.Context()) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond modelPath, _ := createBinFile(t, ggml.KV{ "general.architecture": "llama", "llama.context_length": uint32(32), "llama.embedding_length": uint32(4096), "llama.block_count": uint32(1), "llama.attention.head_count": uint32(32), "llama.attention.head_count_kv": uint32(32), "tokenizer.ggml.tokens": []string{" "}, "tokenizer.ggml.scores": []float32{0}, "tokenizer.ggml.token_type": []int32{0}, }, []*ggml.Tensor{ {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))}, }) reqCtx, cancelReq := context.WithCancel(ctx) defer cancelReq() req := &LlmRequest{ ctx: reqCtx, model: &Model{ModelPath: modelPath}, opts: api.DefaultOptions(), successCh: make(chan *runnerRef, 1), errCh: make(chan error, 1), sessionDuration: &api.Duration{Duration: 2 * time.Minute}, } gpus := []ml.DeviceInfo{} systemInfo := ml.SystemInfo{} server := &mockLlm{vramSize: 10, vramByGPU: map[ml.DeviceID]uint64{}} s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { server.modelPath = model return server, nil } s.load(req, systemInfo, gpus, false) select { case err := <-req.errCh: if err != nil { t.Fatalf("expected no errors when loading, got '%s'", err.Error()) } case resp := <-req.successCh: s.loadedMu.Lock() if resp.refCount != uint(1) || len(s.loaded) != 1 { t.Fatalf("expected a model to be loaded") } s.loadedMu.Unlock() } completedDone := make(chan struct{}) go func() { defer close(completedDone) s.processCompleted(ctx) }() s.expireRunner(&Model{ModelPath: modelPath}) cancelReq() select { case <-s.unloadedCh: case <-time.After(5 * time.Second): t.Fatal("expected model to be unloaded") } s.loadedMu.Lock() if len(s.loaded) != 0 { t.Fatalf("expected model to be unloaded") } s.loadedMu.Unlock() done() select { case <-completedDone: case <-time.After(5 * time.Second): t.Fatal("expected completed loop to stop") } } // TODO - add one scenario that triggers the bogus finished event with positive ref count func TestSchedPrematureExpired(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 1000*time.Millisecond) defer done() // Same model, same request scenario1a := newScenarioRequest(t, ctx, "ollama-model-1a", 10, &api.Duration{Duration: 100 * time.Millisecond}, nil) s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = getGpuFn s.getSystemInfoFn = getSystemInfoFn s.newServerFn = scenario1a.newServer successCh1a, errCh1a := s.GetRunner(scenario1a.ctx, scenario1a.req.model, scenario1a.req.opts, scenario1a.req.sessionDuration) require.Len(t, s.pendingReqCh, 1) s.Run(ctx) select { case resp := <-successCh1a: require.Equal(t, resp.llama, scenario1a.srv) require.Empty(t, s.pendingReqCh) require.Empty(t, errCh1a) s.loadedMu.Lock() require.Len(t, s.loaded, 1) s.loadedMu.Unlock() slog.Info("sending premature expired event now") s.expiredCh <- resp // Shouldn't happen in real life, but make sure its safe case err := <-errCh1a: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } time.Sleep(scenario1a.req.sessionDuration.Duration) scenario1a.ctxDone() time.Sleep(20 * time.Millisecond) require.LessOrEqual(t, len(s.finishedReqCh), 1) time.Sleep(10 * time.Millisecond) require.Empty(t, s.finishedReqCh) s.loadedMu.Lock() require.Empty(t, s.loaded) s.loadedMu.Unlock() // also shouldn't happen in real life s.finishedReqCh <- scenario1a.req time.Sleep(5 * time.Millisecond) } func TestSchedUseLoadedRunner(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) req := &LlmRequest{ ctx: ctx, opts: api.DefaultOptions(), successCh: make(chan *runnerRef, 1), sessionDuration: &api.Duration{Duration: 2}, } finished := make(chan *LlmRequest) llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} r1 := &runnerRef{llama: llm1, sessionDuration: 1, numParallel: 1} req.useLoadedRunner(r1, finished) require.Equal(t, uint(1), r1.refCount) require.Equal(t, time.Duration(2), r1.sessionDuration) select { case success := <-req.successCh: require.Equal(t, r1, success) case err := <-req.errCh: t.Fatal(err.Error()) case <-ctx.Done(): t.Fatal("timeout") } done() fin := <-finished require.Equal(t, req, fin) } func TestSchedUpdateFreeSpace(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() gpus := []ml.DeviceInfo{ { DeviceID: ml.DeviceID{ ID: "1", }, }, { DeviceID: ml.DeviceID{ ID: "2", }, }, } gpus[0].TotalMemory = 1000 gpus[0].FreeMemory = 900 gpus[1].TotalMemory = 2000 gpus[1].FreeMemory = 1900 gpuIDs := []ml.DeviceID{ { ID: "1", }, { ID: "2", }, } llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{{ID: "1"}: 50, {ID: "2"}: 50}} llm2 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{{ID: "1"}: 125, {ID: "2"}: 75}} r1 := &runnerRef{llama: llm1, gpus: gpuIDs, numParallel: 1} r2 := &runnerRef{llama: llm2, gpus: gpuIDs, numParallel: 1} s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.loadedMu.Lock() s.loaded["a"] = r1 s.loaded["b"] = r2 s.loadedMu.Unlock() s.updateFreeSpace(gpus) require.Equal(t, uint64(1000-50-125), gpus[0].FreeMemory) require.Equal(t, uint64(2000-50-75), gpus[1].FreeMemory) } func TestSchedFindRunnerToUnload(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() r1 := &runnerRef{refCount: 1, sessionDuration: 1, numParallel: 1} r2 := &runnerRef{sessionDuration: 2, numParallel: 1} s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.loadedMu.Lock() s.loaded["a"] = r1 s.loaded["b"] = r2 s.loadedMu.Unlock() resp := s.findRunnerToUnload() require.Equal(t, r2, resp) r2.refCount = 1 resp = s.findRunnerToUnload() require.Equal(t, r1, resp) } func TestSchedNeedsReload(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} do := api.DefaultOptions() runner := &runnerRef{ model: &Model{ AdapterPaths: []string{"adapter1"}, ProjectorPaths: []string{"projector1"}, }, Options: &do, llama: llm, numParallel: 1, } req := &LlmRequest{ model: &Model{ AdapterPaths: []string{"adapter2"}, ProjectorPaths: []string{"projector2"}, }, opts: api.DefaultOptions(), } resp := runner.needsReload(ctx, req) require.True(t, resp) req.model.AdapterPaths = runner.model.AdapterPaths resp = runner.needsReload(ctx, req) require.True(t, resp) req.model.ProjectorPaths = runner.model.ProjectorPaths runner.loading = true req.opts.NumBatch = 1234 resp = runner.needsReload(ctx, req) require.True(t, resp) req.opts.NumBatch = runner.Options.NumBatch llm.pingResp = errors.New("foo") resp = runner.needsReload(ctx, req) require.True(t, resp) llm.pingResp = nil resp = runner.needsReload(ctx, req) require.False(t, resp) req.opts.NumGPU = 99 resp = runner.needsReload(ctx, req) require.True(t, resp) req.opts.NumGPU = -1 resp = runner.needsReload(ctx, req) require.False(t, resp) req.contextShift = true resp = runner.needsReload(ctx, req) require.True(t, resp) } func TestResolveContextShift(t *testing.T) { trueValue := true falseValue := false tests := []struct { name string shift *bool model *Model want bool }{ {name: "unset defaults to shift", want: true}, {name: "unset deepseek2 disables shift", model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false}, {name: "unset deepseek2 family disables shift", model: &Model{Config: model.ConfigV2{ModelFamilies: []string{"llama", "deepseek2"}}}, want: false}, {name: "explicit false disables shift", shift: &falseValue, want: false}, {name: "explicit false disables shift for deepseek2", shift: &falseValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false}, {name: "explicit true enables shift", shift: &trueValue, want: true}, {name: "explicit true enables shift for deepseek2", shift: &trueValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.model)) }) } } func TestSchedNeedsReloadIgnoresAutomaticNumCtxClamp(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} opts := api.DefaultOptions() opts.NumCtx = 131072 model := &Model{} runner := &runnerRef{ model: model, Options: &opts, llama: llm, numParallel: 1, numCtxAuto: true, } req := &LlmRequest{ model: model, opts: api.DefaultOptions(), numCtxAuto: true, } req.opts.NumCtx = 262144 require.False(t, runner.needsReload(ctx, req)) req.numCtxAuto = false require.True(t, runner.needsReload(ctx, req)) } func TestSchedNeedsReloadUsesEffectiveAutomaticContextShift(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} opts := api.DefaultOptions() opts.NumCtx = 128 model := &Model{ModelPath: "model.gguf"} runner := &runnerRef{ model: model, Options: &opts, llama: llm, numParallel: 1, numCtxAuto: true, contextShift: true, } req := &LlmRequest{ model: model, opts: api.DefaultOptions(), numCtxAuto: true, } req.opts.NumCtx = 262144 require.False(t, runner.needsReload(ctx, req)) req.numCtxAuto = false require.True(t, runner.needsReload(ctx, req)) } func TestSchedNeedsReloadUsesEffectiveExplicitContext(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} opts := api.DefaultOptions() opts.NumCtx = 2048 model := &Model{ModelPath: "model.gguf"} runner := &runnerRef{ model: model, Options: &opts, llama: llm, numParallel: 1, contextShift: true, trainContext: 2048, } req := &LlmRequest{ model: model, opts: api.DefaultOptions(), } req.opts.NumCtx = 262144 require.False(t, runner.needsReload(ctx, req)) req.opts.NumCtx = 1024 require.True(t, runner.needsReload(ctx, req)) } func TestSchedNeedsReloadIgnoresAutomaticNumBatchDerivation(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} opts := api.DefaultOptions() opts.NumBatch = 1024 model := &Model{} runner := &runnerRef{ model: model, Options: &opts, llama: llm, numParallel: 1, numBatchAuto: true, } req := &LlmRequest{ model: model, opts: api.DefaultOptions(), numBatchAuto: true, } req.opts.NumBatch = 512 require.False(t, runner.needsReload(ctx, req)) req.numBatchAuto = false require.True(t, runner.needsReload(ctx, req)) } func TestSchedNeedsReloadIgnoresAutomaticUseMMapDefault(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} useMmap := false opts := api.DefaultOptions() opts.UseMMap = &useMmap model := &Model{} runner := &runnerRef{ model: model, Options: &opts, llama: llm, numParallel: 1, useMMapAuto: true, } req := &LlmRequest{ model: model, opts: api.DefaultOptions(), } require.False(t, runner.needsReload(ctx, req)) explicitUseMmap := true req.opts.UseMMap = &explicitUseMmap require.True(t, runner.needsReload(ctx, req)) req.opts.UseMMap = &useMmap require.False(t, runner.needsReload(ctx, req)) runner.useMMapAuto = false req.opts.UseMMap = nil require.True(t, runner.needsReload(ctx, req)) } func TestAutomaticGenerationBatch(t *testing.T) { tests := []struct { name string effectiveCtx int predicted uint64 available uint64 flash ml.FlashAttentionType gpus []ml.DeviceInfo want int }{ { name: "small context keeps default", effectiveCtx: 4096, flash: ml.FlashAttentionAuto, want: 512, }, { name: "medium context uses 1024 with unknown memory", effectiveCtx: 32768, flash: ml.FlashAttentionAuto, want: 1024, }, { name: "large context uses 2048 with headroom", effectiveCtx: 131072, predicted: 8 * format.GibiByte, available: 14 * format.GibiByte, flash: ml.FlashAttentionAuto, want: 2048, }, { name: "large context steps down to 1024 without 2048 headroom", effectiveCtx: 131072, predicted: 9 * format.GibiByte, available: 14 * format.GibiByte, flash: ml.FlashAttentionAuto, want: 1024, }, { name: "large context steps down to 1024 for headroom", effectiveCtx: 131072, predicted: 8 * format.GibiByte, available: 11 * format.GibiByte, flash: ml.FlashAttentionAuto, want: 1024, }, { name: "medium context steps down to 512 for headroom", effectiveCtx: 32768, predicted: 8500 * format.MebiByte, available: 11 * format.GibiByte, flash: ml.FlashAttentionAuto, want: 512, }, { name: "flash attention disabled suppresses promotion", effectiveCtx: 131072, predicted: 8 * format.GibiByte, available: 14 * format.GibiByte, flash: ml.FlashAttentionDisabled, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 14 * format.GibiByte}}, want: 512, }, { name: "constrained CUDA without flash attention uses smaller batch", effectiveCtx: 131072, predicted: 3 * format.GibiByte, available: 6 * format.GibiByte, flash: ml.FlashAttentionDisabled, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 6 * format.GibiByte}}, want: 256, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { require.Equal(t, tt.want, automaticGenerationBatch(tt.effectiveCtx, tt.predicted, tt.available, tt.flash, tt.gpus)) }) } } func TestSchedUnloadAllRunners(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond) defer done() llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} llm2 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.unloadAllRunners() r1 := &runnerRef{llama: llm1, numParallel: 1} r2 := &runnerRef{llama: llm2, numParallel: 1} s.loadedMu.Lock() s.loaded["a"] = r1 s.loaded["b"] = r2 s.loadedMu.Unlock() s.unloadAllRunners() require.True(t, llm1.closeCalled) require.True(t, llm2.closeCalled) } func TestSchedUnload(t *testing.T) { llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}} r1 := &runnerRef{llama: llm1, numParallel: 1} r2 := &runnerRef{model: &Model{AdapterPaths: []string{"A"}}, numParallel: 1} r1.unload() require.True(t, llm1.closeCalled) r2.unload() require.Nil(t, r2.model) } func TestSchedAlreadyCanceled(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() dctx, done2 := context.WithCancel(ctx) done2() scenario1a := newScenarioRequest(t, dctx, "ollama-model-1", 10, &api.Duration{Duration: 0}, nil) s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond slog.Info("scenario1a") s.pendingReqCh <- scenario1a.req require.Len(t, s.pendingReqCh, 1) s.Run(ctx) time.Sleep(5 * time.Millisecond) require.Empty(t, s.pendingReqCh) require.Empty(t, scenario1a.req.errCh) require.Empty(t, scenario1a.req.successCh) } // hasLoadedRunner is a test helper that checks if any runner is loaded. func hasLoadedRunner(s *Scheduler) bool { s.loadedMu.Lock() defer s.loadedMu.Unlock() for _, r := range s.loaded { if r.llama != nil { return true } } return false } func TestSchedLlamaServerEvictsWhenVRAMInsufficient(t *testing.T) { // When a llama-server model is predicted to exceed available VRAM, // the scheduler should signal eviction before spawning ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond // GPU with very little free memory — the model won't fit s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 0 // no free VRAM — forces eviction return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn // Pre-load a regular model s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() // Create a request — the model file + KV cache will exceed 100 MiB scenario := newScenarioRequest(t, ctx, "llama-server-model", 1*format.GigaByte, nil, nil) s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return &mockLlm{modelPath: model}, nil } systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.True(t, needEvict, "expected eviction when predicted VRAM exceeds free memory") } func TestSchedLlamaServerExplicitPartialNumGPUSkipsFullFitEviction(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 0 return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() scenario := newScenarioRequest(t, ctx, "partial-llama-server-model", 1*format.GigaByte, nil, nil) scenario.req.opts.NumGPU = 1 scenario.srv.vramSize = 0 called := false s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { called = true require.Equal(t, 1, opts.NumGPU) return scenario.srv, nil } systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "explicit partial offload should not trigger full-fit eviction") require.True(t, called, "scheduler should try the explicitly partial load") } func TestSchedLlamaServerFitsAlongside(t *testing.T) { // When a llama-server model is predicted to fit in remaining VRAM, // it should load without evicting existing models ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond // GPU with plenty of free memory s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn // Pre-load a regular model s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() // The test model GGUF is tiny (~64 bytes) — should easily fit in 20 GiB scenario := newScenarioRequest(t, ctx, "small-llama-server", 1*format.GigaByte, nil, nil) s.newServerFn = scenario.newServer systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) // Should NOT evict — model fits alongside existing needEvict := s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "expected no eviction when model fits in available VRAM") } func TestSchedLlamaServerPredictionUsesTotalParallelContext(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() t.Setenv("OLLAMA_NUM_PARALLEL", "2") s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 900 * format.MebiByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() scenario := newScenarioRequestWithContext(t, ctx, "parallel-context-model", 1*format.GigaByte, nil, nil, 65536) scenario.req.opts.NumCtx = 32768 called := false s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { called = true return scenario.srv, nil } systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.True(t, needEvict, "expected eviction when total parallel context exceeds available memory") require.False(t, called, "preflight prediction should reject before spawning llama-server") } func TestAvailableMemoryForLoadUsesWorstSharedMemoryMeasurement(t *testing.T) { tests := []struct { name string systemFree uint64 gpus []ml.DeviceInfo wantAvailable uint64 wantGPUFree uint64 wantSystemLimited bool }{ { name: "integrated metal uses lower system free", systemFree: 80 * format.GigaByte, gpus: []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "Metal"}, Integrated: true, FreeMemory: 300 * format.GigaByte, }}, wantAvailable: 80 * format.GigaByte, wantGPUFree: 300 * format.GigaByte, wantSystemLimited: true, }, { name: "integrated gpu uses lower system free", systemFree: 6 * format.GigaByte, gpus: []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "Vulkan"}, Integrated: true, FreeMemory: 12 * format.GigaByte, }}, wantAvailable: 6 * format.GigaByte, wantGPUFree: 12 * format.GigaByte, wantSystemLimited: true, }, { name: "discrete metal ignores lower system free", systemFree: 6 * format.GigaByte, gpus: []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "Metal"}, FreeMemory: 12 * format.GigaByte, }}, wantAvailable: 12 * format.GigaByte, wantGPUFree: 12 * format.GigaByte, }, { name: "discrete gpu ignores lower system free", systemFree: 6 * format.GigaByte, gpus: []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 12 * format.GigaByte, }}, wantAvailable: 12 * format.GigaByte, wantGPUFree: 12 * format.GigaByte, }, { name: "mixed gpus only clamp integrated contribution", systemFree: 6 * format.GigaByte, gpus: []ml.DeviceInfo{ { DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 12 * format.GigaByte, }, { DeviceID: ml.DeviceID{Library: "Vulkan"}, Integrated: true, FreeMemory: 10 * format.GigaByte, }, }, wantAvailable: 18 * format.GigaByte, wantGPUFree: 22 * format.GigaByte, wantSystemLimited: true, }, { name: "shared gpu keeps lower adjusted gpu baseline", systemFree: 20 * format.GigaByte, gpus: []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "Metal"}, Integrated: true, FreeMemory: 12 * format.GigaByte, }}, wantAvailable: 12 * format.GigaByte, wantGPUFree: 12 * format.GigaByte, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { available, gpuFree, systemLimited := availableMemoryForLoad(ml.SystemInfo{FreeMemory: tt.systemFree}, tt.gpus) require.Equal(t, tt.wantAvailable, available) require.Equal(t, tt.wantGPUFree, gpuFree) require.Equal(t, tt.wantSystemLimited, systemLimited) }) } } func TestSelectLlamaServerPlacement(t *testing.T) { systemInfo := ml.SystemInfo{FreeMemory: 14 * format.GigaByte} tests := []struct { name string gpus []ml.DeviceInfo predictedVRAM uint64 opts api.Options schedSpread string wantLibrary string wantMainGPU *int wantSelectedGPUs int wantGPUID string }{ { name: "compacts onto largest same-backend GPU", predictedVRAM: 8 * format.GigaByte, gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, Name: "small", FreeMemory: 10 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, Name: "large", FreeMemory: 20 * format.GigaByte}, }, opts: api.DefaultOptions(), wantLibrary: "CUDA", wantMainGPU: testIntPtr(0), wantSelectedGPUs: 1, wantGPUID: "1", }, { name: "explicit main gpu selects matching backend group", predictedVRAM: 8 * format.GigaByte, gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "0", Library: "ROCm"}, FreeMemory: 20 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "ROCm"}, FreeMemory: 24 * format.GigaByte}, }, opts: api.Options{ Runner: api.Runner{MainGPU: testIntPtr(1), NumGPU: -1}, }, wantLibrary: "ROCm", wantMainGPU: testIntPtr(0), wantSelectedGPUs: 1, wantGPUID: "1", }, { name: "integrated GPU is capped by system free memory", predictedVRAM: 12 * format.GigaByte, gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Metal"}, Integrated: true, FreeMemory: 32 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "Metal"}, FreeMemory: 16 * format.GigaByte}, }, opts: api.DefaultOptions(), wantLibrary: "Metal", wantMainGPU: testIntPtr(0), wantSelectedGPUs: 1, wantGPUID: "1", }, { name: "prefers discrete GPU over integrated GPU with more available memory", predictedVRAM: 8 * format.GigaByte, gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Name: "integrated", Integrated: true, FreeMemory: 32 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Name: "discrete", FreeMemory: 10 * format.GigaByte}, }, opts: api.DefaultOptions(), wantLibrary: "Vulkan", wantMainGPU: testIntPtr(0), wantSelectedGPUs: 1, wantGPUID: "1", }, { name: "spread disables automatic compaction", predictedVRAM: 8 * format.GigaByte, schedSpread: "1", gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, FreeMemory: 20 * format.GigaByte}, }, opts: api.DefaultOptions(), wantLibrary: "CUDA", wantSelectedGPUs: 2, }, { name: "no single fit chooses best backend group for llama-server split", predictedVRAM: 30 * format.GigaByte, gpus: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, FreeMemory: 18 * format.GigaByte}, {DeviceID: ml.DeviceID{ID: "0", Library: "ROCm"}, FreeMemory: 12 * format.GigaByte}, }, opts: api.DefaultOptions(), wantLibrary: "CUDA", wantSelectedGPUs: 2, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Setenv("OLLAMA_SCHED_SPREAD", tt.schedSpread) selected, launchOpts := selectLlamaServerPlacement(systemInfo, tt.gpus, tt.predictedVRAM, tt.opts) require.Len(t, selected, tt.wantSelectedGPUs) require.Equal(t, tt.wantLibrary, selected[0].Library) if tt.wantGPUID != "" { require.Equal(t, tt.wantGPUID, selected[0].ID) } if tt.wantMainGPU == nil { require.Nil(t, launchOpts.MainGPU) } else { require.NotNil(t, launchOpts.MainGPU) require.Equal(t, *tt.wantMainGPU, *launchOpts.MainGPU) } }) } } func testIntPtr(v int) *int { return &v } func TestDisableMmapDefaultReason(t *testing.T) { useMmap := true tests := []struct { name string goos string opts api.Options gpus []ml.DeviceInfo blockCount uint64 predictedVRAM uint64 availableVRAM uint64 want string }{ { name: "explicit use_mmap true wins", goos: "windows", opts: api.Options{Runner: api.Runner{NumGPU: -1, UseMMap: &useMmap}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}}, }, { name: "cpu-only request disables mmap", goos: "linux", opts: api.Options{Runner: api.Runner{NumGPU: 0}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}}, want: "cpu", }, { name: "no GPU devices disables mmap", goos: "linux", opts: api.Options{Runner: api.Runner{NumGPU: -1}}, want: "cpu", }, { name: "windows cuda disables mmap", goos: "windows", opts: api.Options{Runner: api.Runner{NumGPU: -1}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}}, want: "windows_cuda", }, { name: "metal partial offload disables mmap", goos: "darwin", opts: api.Options{Runner: api.Runner{NumGPU: 10}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}}, blockCount: 20, want: "metal_partial_offload", }, { name: "metal full offload keeps default", goos: "darwin", opts: api.Options{Runner: api.Runner{NumGPU: 21}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}}, blockCount: 20, }, { name: "metal auto partial offload disables mmap", goos: "darwin", opts: api.Options{Runner: api.Runner{NumGPU: -1}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}}, predictedVRAM: 30 * format.GigaByte, availableVRAM: 20 * format.GigaByte, want: "metal_partial_offload", }, { name: "metal auto full offload keeps default", goos: "darwin", opts: api.Options{Runner: api.Runner{NumGPU: -1}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}}, predictedVRAM: 10 * format.GigaByte, availableVRAM: 20 * format.GigaByte, }, { name: "linux cuda keeps default", goos: "linux", opts: api.Options{Runner: api.Runner{NumGPU: -1}}, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}}, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { require.Equal(t, tt.want, disableMmapDefaultReason(tt.goos, tt.opts, tt.gpus, tt.blockCount, tt.predictedVRAM, tt.availableVRAM)) }) } } func TestDisableMmapForHostPressure(t *testing.T) { gpus := []ml.DeviceInfo{{ DeviceID: ml.DeviceID{Library: "CUDA"}, TotalMemory: 100 * format.GigaByte, FreeMemory: 80 * format.GigaByte, }} systemInfo := ml.SystemInfo{ TotalMemory: 100 * format.GigaByte, FreeMemory: 50 * format.GigaByte, } require.True(t, disableMmapForHostPressure( "linux", api.Options{}, systemInfo, gpus, 20*format.GigaByte, 25*format.GigaByte, 30*format.GigaByte, 80*format.GigaByte, )) useMmap := true require.False(t, disableMmapForHostPressure( "linux", api.Options{Runner: api.Runner{UseMMap: &useMmap}}, systemInfo, gpus, 20*format.GigaByte, 25*format.GigaByte, 30*format.GigaByte, 80*format.GigaByte, ), "explicit use_mmap=true should win") require.False(t, disableMmapForHostPressure( "darwin", api.Options{}, systemInfo, gpus, 20*format.GigaByte, 25*format.GigaByte, 30*format.GigaByte, 80*format.GigaByte, ), "only the Linux pressure heuristic is restored") igpu := append([]ml.DeviceInfo(nil), gpus...) igpu[0].Integrated = true require.False(t, disableMmapForHostPressure( "linux", api.Options{}, systemInfo, igpu, 20*format.GigaByte, 25*format.GigaByte, 30*format.GigaByte, 80*format.GigaByte, ), "shared-memory GPU loads should keep the normal mmap path") require.False(t, disableMmapForHostPressure( "linux", api.Options{}, systemInfo, gpus, 20*format.GigaByte, 25*format.GigaByte, 70*format.GigaByte, 80*format.GigaByte, ), "when VRAM is tight, no-mmap could make partial CPU offload worse") } // TestSchedLoadCrashTriggersEvictAllAndRetry verifies that a post-spawn // Load() OOM while other models are resident signals evict-all-and-retry // on the first attempt, but fails fast on the second attempt. func TestSchedLoadCrashTriggersEvictAllAndRetry(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn // Pre-load a different model so evict-all has something to evict. s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "/fake/existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() // newServerFn returns a mockLlm that crashes in Load() loadCrash := errors.New("cudaMalloc failed: out of memory") s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil) systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) // First attempt: should signal evict-all by returning true and NOT send // the error to errCh (so the caller will retry). needEvict := s.load(scenario.req, systemInfo, gpus, true) require.True(t, needEvict, "first load OOM should signal eviction") require.True(t, scenario.req.oomRetryAttempted, "oomRetryAttempted should be set") select { case err := <-scenario.req.errCh: t.Fatalf("errCh should be empty on first crash, got %v", err) default: } // Second attempt (simulating the retry after processPending evicted all // other runners): same crash, but this time oomRetryAttempted is set so // load() should fail fast and report the error. needEvict = s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "second load OOM should not ask for another eviction") select { case err := <-scenario.req.errCh: require.ErrorIs(t, err, loadCrash) case <-time.After(100 * time.Millisecond): t.Fatal("expected error on errCh after second crash") } } func TestSchedLoadOOMReducesAutomaticContextBeforeRetry(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "/fake/existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() loadCrash := errors.New("cudaMalloc failed: out of memory") var seenNumCtx []int var seenNumBatch []int s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { seenNumCtx = append(seenNumCtx, opts.NumCtx) seenNumBatch = append(seenNumBatch, opts.NumBatch) return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } scenario := newScenarioRequestWithContext(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil, 131072) scenario.req.opts.NumCtx = 262144 scenario.req.numCtxAuto = true scenario.req.numBatchAuto = true systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.True(t, needEvict, "first automatic-context load OOM should signal eviction and retry") require.True(t, scenario.req.oomRetryAttempted) require.Equal(t, 32768, scenario.req.opts.NumCtx) require.Equal(t, 1024, scenario.req.opts.NumBatch) select { case err := <-scenario.req.errCh: t.Fatalf("errCh should be empty on first crash, got %v", err) default: } needEvict = s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "second load OOM should not ask for another eviction") require.Equal(t, []int{262144, 32768}, seenNumCtx) require.Equal(t, []int{2048, 1024}, seenNumBatch) select { case err := <-scenario.req.errCh: require.ErrorIs(t, err, loadCrash) case <-time.After(100 * time.Millisecond): t.Fatal("expected error on errCh after second crash") } } func TestSchedLoadOOMKeepsExplicitContextBeforeRetry(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "/fake/existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() loadCrash := errors.New("cudaMalloc failed: out of memory") s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } scenario := newScenarioRequestWithContext(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil, 131072) scenario.req.opts.NumCtx = 262144 scenario.req.numCtxAuto = false systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.True(t, needEvict, "explicit-context load OOM should still evict and retry once") require.True(t, scenario.req.oomRetryAttempted) require.Equal(t, 262144, scenario.req.opts.NumCtx) select { case err := <-scenario.req.errCh: t.Fatalf("errCh should be empty on first crash, got %v", err) default: } } func TestSchedFirstLoadOOMReducesAutomaticContextAndRetries(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), time.Second) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn loadCrash := errors.New("cudaMalloc failed: out of memory") var seenNumCtx []int s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { seenNumCtx = append(seenNumCtx, opts.NumCtx) if len(seenNumCtx) == 1 { return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } return &mockLlm{modelPath: model, vramSize: 1 * format.GigaByte, contextLength: opts.NumCtx}, nil } scenario := newScenarioRequestWithContext(t, ctx, "first-load-crashing-model", 1*format.GigaByte, nil, nil, 131072) scenario.req.opts.NumCtx = 262144 scenario.req.numCtxAuto = true s.pendingReqCh <- scenario.req s.Run(ctx) select { case runner := <-scenario.req.successCh: require.Equal(t, 32768, runner.Options.NumCtx) require.Equal(t, []int{262144, 32768}, seenNumCtx) case err := <-scenario.req.errCh: t.Fatalf("expected retry success, got error %v", err) case <-ctx.Done(): t.Fatal("timed out waiting for first-load retry") } } // TestSchedLoadCrashNoOtherModelsFailsFast verifies that a Load() crash with // no other resident models reports the error immediately (no retry). func TestSchedLoadCrashNoOtherModelsFailsFast(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn loadCrash := errors.New("simulated llama-server OOM crash") s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil) systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "crash with no other runners should not ask for eviction") require.False(t, scenario.req.oomRetryAttempted, "oomRetryAttempted must stay false") select { case err := <-scenario.req.errCh: require.ErrorIs(t, err, loadCrash) case <-time.After(100 * time.Millisecond): t.Fatal("expected error on errCh immediately") } } func TestSchedLoadNonOOMWithOtherModelsFailsFast(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.waitForRecovery = 10 * time.Millisecond s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo { g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}} g.TotalMemory = 24 * format.GigaByte g.FreeMemory = 20 * format.GigaByte return []ml.DeviceInfo{g} } s.getSystemInfoFn = getSystemInfoFn s.loadedMu.Lock() s.loaded["existing-model"] = &runnerRef{ llama: &mockLlm{modelPath: "/fake/existing"}, modelKey: "existing-model", } s.loadedMu.Unlock() loadCrash := errors.New("server parse failed") s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) { return &mockLlm{modelPath: model, loadErr: loadCrash}, nil } scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil) systemInfo := getSystemInfoFn() gpus := s.getGpuFn(ctx, nil) needEvict := s.load(scenario.req, systemInfo, gpus, true) require.False(t, needEvict, "non-OOM load crash should not ask for eviction") require.False(t, scenario.req.oomRetryAttempted, "oomRetryAttempted must stay false") select { case err := <-scenario.req.errCh: require.ErrorIs(t, err, loadCrash) case <-time.After(100 * time.Millisecond): t.Fatal("expected error on errCh immediately") } } func TestSchedRuntimeOOMExpiresLoadedRunners(t *testing.T) { ctx, done := context.WithCancel(t.Context()) defer done() s := InitScheduler(ctx) currentModel := &Model{ModelPath: "/tmp/current.gguf"} current := &runnerRef{ model: currentModel, modelKey: schedulerModelKey(currentModel), sessionDuration: time.Hour, llama: &mockLlm{modelPath: "/tmp/current.gguf"}, } otherModel := &Model{ModelPath: "/tmp/other.gguf"} other := &runnerRef{ model: otherModel, modelKey: schedulerModelKey(otherModel), sessionDuration: time.Hour, llama: &mockLlm{modelPath: "/tmp/other.gguf"}, } s.loadedMu.Lock() s.loaded[current.modelKey] = current s.loaded[other.modelKey] = other s.loadedMu.Unlock() s.expireRunnersForRuntimeOOM(currentModel, errors.New("cudaMalloc failed: out of memory")) require.Equal(t, time.Duration(0), current.sessionDuration) require.Equal(t, time.Duration(0), other.sessionDuration) require.Len(t, s.expiredCh, 2) } func TestSchedLlamaServerEvictsExistingOnPending(t *testing.T) { // When a llama-server runner is already loaded and a new model is requested, // the scheduler should evict the llama-server runner ctx, done := context.WithCancel(t.Context()) defer done() s := InitScheduler(ctx) // Load a llama-server runner s.loadedMu.Lock() s.loaded["llama-model"] = &runnerRef{ llama: &mockLlm{modelPath: "/tmp/model.gguf"}, modelKey: "llama-model", } s.loadedMu.Unlock() require.True(t, hasLoadedRunner(s)) // The findRunnerToUnload should find and return the llama-server runner runner := s.findRunnerToUnload() require.NotNil(t, runner) } type mockLlm struct { modelPath string pingResp error waitResp error completionResp error embeddingResp []float32 embeddingRespErr error tokenizeResp []int tokenizeRespErr error detokenizeResp string detonekizeRespErr error closeResp error closeCalled bool vramSize uint64 totalSize uint64 contextLength int vramByGPU map[ml.DeviceID]uint64 // loadErr, if non-nil, is returned from Load() to simulate a post-spawn // load failure (e.g. llama-server crashing due to under-predicted VRAM). loadErr error } func (s *mockLlm) ModelPath() string { return s.modelPath } func (s *mockLlm) Load(ctx context.Context, sytemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) { if s.loadErr != nil { return nil, s.loadErr } if requireFull { if len(gpus) == 0 { slog.Info("mockLlm.Load CPU based load") return nil, nil } for _, g := range gpus { if g.FreeMemory >= s.vramSize { return []ml.DeviceID{g.DeviceID}, nil } } return nil, llm.ErrLoadRequiredFull } gpuIDs := make([]ml.DeviceID, len(gpus)) for i := range gpus { gpuIDs[i] = gpus[i].DeviceID } return gpuIDs, nil } func (s *mockLlm) Ping(ctx context.Context) error { return s.pingResp } func (s *mockLlm) WaitUntilRunning(ctx context.Context) error { return s.waitResp } func (s *mockLlm) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error { return s.completionResp } func (s *mockLlm) Chat(ctx context.Context, req llm.ChatRequest, fn func(llm.ChatResponse)) error { return errors.New("not implemented") } func (s *mockLlm) ApplyChatTemplate(ctx context.Context, req llm.ChatRequest) (string, error) { return "", errors.New("not implemented") } func (s *mockLlm) Embedding(ctx context.Context, input string) ([]float32, int, error) { return s.embeddingResp, 0, s.embeddingRespErr } func (s *mockLlm) Tokenize(ctx context.Context, content string) ([]int, error) { return s.tokenizeResp, s.tokenizeRespErr } func (s *mockLlm) Detokenize(ctx context.Context, tokens []int) (string, error) { return s.detokenizeResp, s.detonekizeRespErr } func (s *mockLlm) Close() error { s.closeCalled = true return s.closeResp } func (s *mockLlm) MemorySize() (uint64, uint64) { return s.totalSize, s.vramSize } func (s *mockLlm) VRAMByGPU(id ml.DeviceID) uint64 { return s.vramByGPU[id] } func (s *mockLlm) Pid() int { return -1 } func (s *mockLlm) GetPort() int { return -1 } func (s *mockLlm) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo { return nil } func (s *mockLlm) HasExited() bool { return false } func (s *mockLlm) GetActiveDeviceIDs() []ml.DeviceID { return nil } func (s *mockLlm) ContextLength() int { return s.contextLength } // TestImageGenRunnerCanBeEvicted verifies that an image generation model // loaded in the scheduler can be evicted when idle. func TestImageGenRunnerCanBeEvicted(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.getGpuFn = getGpuFn s.getSystemInfoFn = getSystemInfoFn // Simulate an image gen runner already loaded imageGenRunner := &runnerRef{ model: &Model{Name: "z-image", ModelPath: "/fake/image/model"}, modelPath: "/fake/image/model", llama: &mockLlm{vramSize: 21 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{}}, sessionDuration: 5 * time.Millisecond, refCount: 0, // idle } s.loadedMu.Lock() s.loaded["/fake/image/model"] = imageGenRunner s.loadedMu.Unlock() // Verify the image gen runner is loaded s.loadedMu.Lock() require.Len(t, s.loaded, 1) s.loadedMu.Unlock() // findRunnerToUnload should find the idle image gen runner runner := s.findRunnerToUnload() require.NotNil(t, runner) require.Equal(t, "/fake/image/model", runner.modelPath) } // TestImageGenSchedulerCoexistence verifies that image generation models // can coexist with language models in the scheduler and VRAM is tracked correctly. func TestImageGenSchedulerCoexistence(t *testing.T) { ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond) defer done() s := InitScheduler(ctx) s.getGpuFn = getGpuFn s.getSystemInfoFn = getSystemInfoFn // Load both an imagegen runner and a language model runner imageGenRunner := &runnerRef{ model: &Model{Name: "flux", ModelPath: "/fake/flux/model"}, modelPath: "/fake/flux/model", llama: &mockLlm{vramSize: 8 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 8 * format.GigaByte}}, sessionDuration: 10 * time.Millisecond, numParallel: 1, refCount: 0, } langModelRunner := &runnerRef{ model: &Model{Name: "llama3", ModelPath: "/fake/llama3/model"}, modelPath: "/fake/llama3/model", llama: &mockLlm{vramSize: 4 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 4 * format.GigaByte}}, sessionDuration: 10 * time.Millisecond, numParallel: 1, refCount: 0, } s.loadedMu.Lock() s.loaded["/fake/flux/model"] = imageGenRunner s.loaded["/fake/llama3/model"] = langModelRunner s.loadedMu.Unlock() // Verify both are loaded s.loadedMu.Lock() require.Len(t, s.loaded, 2) require.NotNil(t, s.loaded["/fake/flux/model"]) require.NotNil(t, s.loaded["/fake/llama3/model"]) s.loadedMu.Unlock() // Verify updateFreeSpace accounts for both gpus := []ml.DeviceInfo{ { DeviceID: ml.DeviceID{Library: "Metal"}, TotalMemory: 24 * format.GigaByte, FreeMemory: 24 * format.GigaByte, }, } s.updateFreeSpace(gpus) // Free memory should be reduced by both models expectedFree := uint64(24*format.GigaByte) - uint64(8*format.GigaByte) - uint64(4*format.GigaByte) require.Equal(t, expectedFree, gpus[0].FreeMemory) }