mirror of
https://github.com/ollama/ollama.git
synced 2026-07-10 09:44:18 +00:00
2248 lines
72 KiB
Go
2248 lines
72 KiB
Go
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)
|
|
}
|