sched_test.go 17.5 KB
Newer Older
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
package server

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log/slog"
	"os"
	"testing"
	"time"

	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/app/lifecycle"
Michael Yang's avatar
lint  
Michael Yang committed
15
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/gpu"
	"github.com/ollama/ollama/llm"
	"github.com/stretchr/testify/require"
)

func init() {
	os.Setenv("OLLAMA_DEBUG", "1")
	lifecycle.InitLogging()
}

func TestInitScheduler(t *testing.T) {
	ctx, done := context.WithCancel(context.Background())
	defer done()
	s := InitScheduler(ctx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
31
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
32
	require.NotNil(t, s.loaded)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
33
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
34
35
36
}

func TestLoad(t *testing.T) {
37
	ctx, done := context.WithTimeout(context.Background(), 20*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
38
39
	defer done()
	s := InitScheduler(ctx)
40
	var ggml *llm.GGML // value not used in tests
Daniel Hiltgen's avatar
Daniel Hiltgen committed
41
42
43
	req := &LlmRequest{
		ctx:             ctx,
		model:           &Model{ModelPath: "foo"},
Daniel Hiltgen's avatar
Daniel Hiltgen committed
44
		opts:            api.DefaultOptions(),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
45
46
47
48
49
50
51
52
53
		successCh:       make(chan *runnerRef, 1),
		errCh:           make(chan error, 1),
		sessionDuration: 2,
	}
	// Fail to load model first
	s.newServerFn = func(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
		return nil, fmt.Errorf("something failed to load model blah")
	}
	gpus := gpu.GpuInfoList{}
54
	s.load(req, ggml, gpus)
Michael Yang's avatar
lint  
Michael Yang committed
55
	require.Empty(t, req.successCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
56
	require.Len(t, req.errCh, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
57
	s.loadedMu.Lock()
Michael Yang's avatar
lint  
Michael Yang committed
58
	require.Empty(t, s.loaded)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
59
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
60
61
62
63
64
65
66
	err := <-req.errCh
	require.Contains(t, err.Error(), "this model may be incompatible")

	server := &mockLlm{estimatedVRAM: 10}
	s.newServerFn = func(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
		return server, nil
	}
67
	s.load(req, ggml, gpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
68
69
70
71
72
73
	select {
	case err := <-req.errCh:
		require.NoError(t, err)
	case resp := <-req.successCh:
		require.Equal(t, uint64(10), resp.estimatedVRAM)
		require.Equal(t, uint(1), resp.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
74
		s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
75
		require.Len(t, s.loaded, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
76
		s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
77
78
79
80
	}

	req.model.ModelPath = "dummy_model_path"
	server.waitResp = fmt.Errorf("wait failure")
81
	s.load(req, ggml, gpus)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
82
83
84
85
86
87
	select {
	case err := <-req.errCh:
		require.Contains(t, err.Error(), "wait failure")
	case resp := <-req.successCh:
		t.Errorf("unexpected success %v", resp)
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
88
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
89
	runner := s.loaded["dummy_model_path"]
Daniel Hiltgen's avatar
Daniel Hiltgen committed
90
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
91
92
	require.NotNil(t, runner)
	require.Equal(t, uint(0), runner.refCount)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
93
	time.Sleep(1 * time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
94
95
96
97
98
99
100
101
	require.Len(t, s.expiredCh, 1)
}

type bundle struct {
	ctx     context.Context //nolint:containedctx
	ctxDone func()
	srv     *mockLlm
	req     *LlmRequest
102
	ggml    *llm.GGML
Daniel Hiltgen's avatar
Daniel Hiltgen committed
103
104
105
106
107
108
109
110
111
112
113
114
}

func (scenario *bundle) newServer(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
	return scenario.srv, nil
}

func newScenario(t *testing.T, ctx context.Context, modelName string, estimatedVRAM uint64) *bundle {
	scenario := &bundle{}
	scenario.ctx, scenario.ctxDone = context.WithCancel(ctx)
	t.Helper()

	f, err := os.CreateTemp(t.TempDir(), modelName)
Michael Yang's avatar
lint  
Michael Yang committed
115
	require.NoError(t, err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
	defer f.Close()

	gguf := llm.NewGGUFV3(binary.LittleEndian)
	err = gguf.Encode(f, llm.KV{
		"general.architecture":          "llama",
		"general.name":                  "name",
		"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},
	}, []llm.Tensor{
		{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: &bytes.Reader{}},
132
		{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: &bytes.Reader{}},
Daniel Hiltgen's avatar
Daniel Hiltgen committed
133
	})
Michael Yang's avatar
lint  
Michael Yang committed
134
	require.NoError(t, err)
135

Daniel Hiltgen's avatar
Daniel Hiltgen committed
136
137
	fname := f.Name()
	model := &Model{Name: modelName, ModelPath: fname}
138
	scenario.ggml, err = llm.LoadModel(model.ModelPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
	require.NoError(t, err)
140

Daniel Hiltgen's avatar
Daniel Hiltgen committed
141
142
143
	scenario.req = &LlmRequest{
		ctx:             scenario.ctx,
		model:           model,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
144
		opts:            api.DefaultOptions(),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
145
146
147
148
149
150
151
152
153
		sessionDuration: 5 * time.Millisecond,
		successCh:       make(chan *runnerRef, 1),
		errCh:           make(chan error, 1),
	}
	scenario.srv = &mockLlm{estimatedVRAM: estimatedVRAM}
	return scenario
}

func TestRequests(t *testing.T) {
154
	ctx, done := context.WithTimeout(context.Background(), time.Second)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
155
156
157
158
159
160
161
	defer done()

	// Same model, same request
	scenario1a := newScenario(t, ctx, "ollama-model-1", 10)
	scenario1a.req.sessionDuration = 0
	scenario1b := newScenario(t, ctx, "ollama-model-1", 11)
	scenario1b.req.model = scenario1a.req.model
162
	scenario1b.ggml = scenario1a.ggml
Daniel Hiltgen's avatar
Daniel Hiltgen committed
163
164
165
166
	scenario1b.req.sessionDuration = 0

	// simple reload of same model
	scenario2a := newScenario(t, ctx, "ollama-model-1", 20)
167
168
	tmpModel := *scenario1a.req.model
	scenario2a.req.model = &tmpModel
169
	scenario2a.ggml = scenario1a.ggml
Daniel Hiltgen's avatar
Daniel Hiltgen committed
170
171
172
173

	// Multiple loaded models
	scenario3a := newScenario(t, ctx, "ollama-model-3a", 1*format.GigaByte)
	scenario3b := newScenario(t, ctx, "ollama-model-3b", 24*format.GigaByte)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
174
175
176
	scenario3c := newScenario(t, ctx, "ollama-model-4a", 30)
	scenario3c.req.opts.NumGPU = 0                           // CPU load, will be allowed
	scenario3d := newScenario(t, ctx, "ollama-model-3c", 30) // Needs prior unloaded
Daniel Hiltgen's avatar
Daniel Hiltgen committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192

	s := InitScheduler(ctx)
	s.getGpuFn = func() gpu.GpuInfoList {
		g := gpu.GpuInfo{Library: "metal"}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 12 * format.GigaByte
		return []gpu.GpuInfo{g}
	}
	s.newServerFn = scenario1a.newServer
	slog.Info("scenario1a")
	s.pendingReqCh <- scenario1a.req
	require.Len(t, s.pendingReqCh, 1)
	s.Run(ctx)
	select {
	case resp := <-scenario1a.req.successCh:
		require.Equal(t, resp.llama, scenario1a.srv)
Michael Yang's avatar
lint  
Michael Yang committed
193
194
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario1a.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
195
196
197
198
199
200
201
202
203
204
205
	case <-ctx.Done():
		t.Errorf("timeout")
	}

	// Same runner as first request due to not needing a reload
	s.newServerFn = scenario1b.newServer
	slog.Info("scenario1b")
	s.pendingReqCh <- scenario1b.req
	select {
	case resp := <-scenario1b.req.successCh:
		require.Equal(t, resp.llama, scenario1a.srv)
Michael Yang's avatar
lint  
Michael Yang committed
206
207
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario1b.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
	case <-ctx.Done():
		t.Errorf("timeout")
	}

	// Trigger a reload
	s.newServerFn = scenario2a.newServer
	scenario2a.req.model.AdapterPaths = []string{"new"}
	slog.Info("scenario2a")
	s.pendingReqCh <- scenario2a.req
	// finish first two requests, so model can reload
	time.Sleep(1 * time.Millisecond)
	scenario1a.ctxDone()
	scenario1b.ctxDone()
	select {
	case resp := <-scenario2a.req.successCh:
		require.Equal(t, resp.llama, scenario2a.srv)
Michael Yang's avatar
lint  
Michael Yang committed
224
225
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario2a.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
226
227
228
229
	case <-ctx.Done():
		t.Errorf("timeout")
	}

230
	envconfig.MaxRunners = 1
Daniel Hiltgen's avatar
Daniel Hiltgen committed
231
232
233
234
235
236
237
238
239
	s.newServerFn = scenario3a.newServer
	slog.Info("scenario3a")
	s.pendingReqCh <- scenario3a.req
	// finish prior request, so new model can load
	time.Sleep(1 * time.Millisecond)
	scenario2a.ctxDone()
	select {
	case resp := <-scenario3a.req.successCh:
		require.Equal(t, resp.llama, scenario3a.srv)
Michael Yang's avatar
lint  
Michael Yang committed
240
241
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario3a.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
242
243
244
	case <-ctx.Done():
		t.Errorf("timeout")
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
245
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
246
	require.Len(t, s.loaded, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
247
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
248

249
	envconfig.MaxRunners = 0
Daniel Hiltgen's avatar
Daniel Hiltgen committed
250
251
252
253
254
255
	s.newServerFn = scenario3b.newServer
	slog.Info("scenario3b")
	s.pendingReqCh <- scenario3b.req
	select {
	case resp := <-scenario3b.req.successCh:
		require.Equal(t, resp.llama, scenario3b.srv)
Michael Yang's avatar
lint  
Michael Yang committed
256
257
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario3b.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
258
259
260
	case <-ctx.Done():
		t.Errorf("timeout")
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
261
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
262
	require.Len(t, s.loaded, 2)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
263
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
264

Daniel Hiltgen's avatar
Daniel Hiltgen committed
265
	// This is a CPU load with NumGPU = 0 so it should load
Daniel Hiltgen's avatar
Daniel Hiltgen committed
266
267
	s.newServerFn = scenario3c.newServer
	slog.Info("scenario3c")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
268
269
270
271
	s.pendingReqCh <- scenario3c.req
	select {
	case resp := <-scenario3c.req.successCh:
		require.Equal(t, resp.llama, scenario3c.srv)
Michael Yang's avatar
lint  
Michael Yang committed
272
273
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario3c.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
274
275
276
277
278
279
280
281
282
283
284
285
286
	case <-ctx.Done():
		t.Errorf("timeout")
	}
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 3)
	s.loadedMu.Unlock()

	// Try to load a model that wont fit
	s.newServerFn = scenario3d.newServer
	slog.Info("scenario3d")
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 3)
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
287
288
	scenario3a.ctxDone() // Won't help since this one isn't big enough to make room
	time.Sleep(2 * time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
289
	s.pendingReqCh <- scenario3d.req
Daniel Hiltgen's avatar
Daniel Hiltgen committed
290
291
	// finish prior request, so new model can load
	time.Sleep(6 * time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
292
293
294
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
295
296
	scenario3b.ctxDone()
	select {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
297
298
	case resp := <-scenario3d.req.successCh:
		require.Equal(t, resp.llama, scenario3d.srv)
Michael Yang's avatar
lint  
Michael Yang committed
299
300
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, scenario3d.req.errCh)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
301
302
303
	case <-ctx.Done():
		t.Errorf("timeout")
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
304
305
306
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
307
308
309
}

func TestGetRunner(t *testing.T) {
310
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
311
312
313
314
315
316
317
318
319
	defer done()

	// Same model, same request
	scenario1a := newScenario(t, ctx, "ollama-model-1a", 10)
	scenario1a.req.sessionDuration = 0
	scenario1b := newScenario(t, ctx, "ollama-model-1b", 10)
	scenario1b.req.sessionDuration = 0
	scenario1c := newScenario(t, ctx, "ollama-model-1c", 10)
	scenario1c.req.sessionDuration = 0
320
	envconfig.MaxQueuedRequests = 1
Daniel Hiltgen's avatar
Daniel Hiltgen committed
321
322
323
324
325
326
327
328
329
330
331
332
333
334
	s := InitScheduler(ctx)
	s.getGpuFn = func() gpu.GpuInfoList {
		g := gpu.GpuInfo{Library: "metal"}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 12 * format.GigaByte
		return []gpu.GpuInfo{g}
	}
	s.newServerFn = scenario1a.newServer
	slog.Info("scenario1a")
	successCh1a, errCh1a := s.GetRunner(scenario1a.ctx, scenario1a.req.model, scenario1a.req.opts, scenario1a.req.sessionDuration)
	require.Len(t, s.pendingReqCh, 1)
	slog.Info("scenario1b")
	successCh1b, errCh1b := s.GetRunner(scenario1b.ctx, scenario1b.req.model, scenario1b.req.opts, scenario1b.req.sessionDuration)
	require.Len(t, s.pendingReqCh, 1)
Michael Yang's avatar
lint  
Michael Yang committed
335
	require.Empty(t, successCh1b)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
336
337
338
339
340
341
342
	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, scenario1a.srv)
Michael Yang's avatar
lint  
Michael Yang committed
343
344
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, errCh1a)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
345
346
347
348
	case <-ctx.Done():
		t.Errorf("timeout")
	}
	scenario1a.ctxDone()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
349
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
350
	require.Len(t, s.loaded, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
351
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
352
353
354
355

	scenario1c.req.model.ModelPath = "bad path"
	slog.Info("scenario1c")
	successCh1c, errCh1c := s.GetRunner(scenario1c.ctx, scenario1c.req.model, scenario1c.req.opts, scenario1c.req.sessionDuration)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
356
	// Starts in pending channel, then should be quickly processsed to return an error
357
	time.Sleep(5 * time.Millisecond)
Michael Yang's avatar
lint  
Michael Yang committed
358
	require.Empty(t, successCh1c)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
359
	s.loadedMu.Lock()
Michael Yang's avatar
lint  
Michael Yang committed
360
	require.Empty(t, s.loaded)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
361
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
362
363
364
365
366
367
368
369
	require.Len(t, errCh1c, 1)
	err = <-errCh1c
	require.Contains(t, err.Error(), "bad path")
	scenario1b.ctxDone()
}

// TODO - add one scenario that triggers the bogus finished event with positive ref count
func TestPrematureExpired(t *testing.T) {
370
	ctx, done := context.WithTimeout(context.Background(), 500*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
	defer done()

	// Same model, same request
	scenario1a := newScenario(t, ctx, "ollama-model-1a", 10)
	s := InitScheduler(ctx)
	s.getGpuFn = func() gpu.GpuInfoList {
		g := gpu.GpuInfo{Library: "metal"}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 12 * format.GigaByte
		return []gpu.GpuInfo{g}
	}
	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)
Michael Yang's avatar
lint  
Michael Yang committed
389
390
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, errCh1a)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
391
		s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
392
		require.Len(t, s.loaded, 1)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
393
		s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
394
395
396
397
398
399
400
401
402
403
		slog.Info("sending premature expired event now")
		s.expiredCh <- resp // Shouldn't happen in real life, but make sure its safe
	case <-ctx.Done():
		t.Errorf("timeout")
	}
	time.Sleep(scenario1a.req.sessionDuration)
	scenario1a.ctxDone()
	time.Sleep(20 * time.Millisecond)
	require.LessOrEqual(t, len(s.finishedReqCh), 1)
	time.Sleep(10 * time.Millisecond)
Michael Yang's avatar
lint  
Michael Yang committed
404
	require.Empty(t, s.finishedReqCh)
405
	s.loadedMu.Lock()
Michael Yang's avatar
lint  
Michael Yang committed
406
	require.Empty(t, s.loaded)
407
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
408
409
410
411
412
413
414

	// also shouldn't happen in real life
	s.finishedReqCh <- scenario1a.req
	time.Sleep(5 * time.Millisecond)
}

func TestUseLoadedRunner(t *testing.T) {
415
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
416
417
	req := &LlmRequest{
		ctx:             ctx,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
418
		opts:            api.DefaultOptions(),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
		successCh:       make(chan *runnerRef, 1),
		sessionDuration: 2,
	}
	finished := make(chan *LlmRequest)
	llm1 := &mockLlm{}
	r1 := &runnerRef{llama: llm1, sessionDuration: 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 <-ctx.Done():
		t.Errorf("timeout")
	}
	done()
	fin := <-finished
	require.Equal(t, req, fin)
}

func TestUpdateFreeSpace(t *testing.T) {
440
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
	defer done()
	gpus := gpu.GpuInfoList{
		{
			Library: "a",
			ID:      "1",
		},
		{
			Library: "a",
			ID:      "2",
		},
	}
	gpus[0].TotalMemory = 1000
	gpus[0].FreeMemory = 900
	gpus[1].TotalMemory = 2000
	gpus[1].FreeMemory = 1900
	llm1 := &mockLlm{estimatedVRAM: 100}
	llm2 := &mockLlm{estimatedVRAM: 200}
	r1 := &runnerRef{llama: llm1, gpus: gpus}
	r2 := &runnerRef{llama: llm2, gpus: gpus}

	s := InitScheduler(ctx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
462
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
463
464
	s.loaded["a"] = r1
	s.loaded["b"] = r2
Daniel Hiltgen's avatar
Daniel Hiltgen committed
465
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
466
467
468
469
470
471
472

	s.updateFreeSpace(gpus)
	require.Equal(t, uint64(850), gpus[0].FreeMemory)
	require.Equal(t, uint64(1850), gpus[1].FreeMemory)
}

func TestFindRunnerToUnload(t *testing.T) {
473
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
474
	defer done()
475

Daniel Hiltgen's avatar
Daniel Hiltgen committed
476
477
478
479
	r1 := &runnerRef{refCount: 1, sessionDuration: 1}
	r2 := &runnerRef{sessionDuration: 2}

	s := InitScheduler(ctx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
480
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
481
482
	s.loaded["a"] = r1
	s.loaded["b"] = r2
Daniel Hiltgen's avatar
Daniel Hiltgen committed
483
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
484

485
	resp := s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
486
487
	require.Equal(t, r2, resp)
	r2.refCount = 1
488
	resp = s.findRunnerToUnload()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
489
490
491
492
	require.Equal(t, r1, resp)
}

func TestNeedsReload(t *testing.T) {
493
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
494
495
496
	defer done()

	llm := &mockLlm{}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
497
	do := api.DefaultOptions()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
498
	runner := &runnerRef{
499
500
501
		model:   &Model{AdapterPaths: []string{"adapter1"}, ProjectorPaths: []string{"projector1"}},
		Options: &do,
		llama:   llm,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
502
503
504
505
506
507
	}
	req := &LlmRequest{
		model: &Model{
			AdapterPaths:   []string{"adapter2"},
			ProjectorPaths: []string{"projector2"},
		},
Daniel Hiltgen's avatar
Daniel Hiltgen committed
508
		opts: api.DefaultOptions(),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
509
510
511
	}
	resp := runner.needsReload(ctx, req)
	require.True(t, resp)
512
	req.model.AdapterPaths = runner.model.AdapterPaths
Daniel Hiltgen's avatar
Daniel Hiltgen committed
513
514
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
515
	req.model.ProjectorPaths = runner.model.ProjectorPaths
Daniel Hiltgen's avatar
Daniel Hiltgen committed
516
517
518
519
520
521
522
523
524
525
526
527
528
	runner.loading = true
	req.opts.NumBatch = 1234
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
	req.opts.NumBatch = runner.Options.NumBatch
	llm.pingResp = fmt.Errorf("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)
529
530
531
	require.True(t, resp)
	req.opts.NumGPU = -1
	resp = runner.needsReload(ctx, req)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
532
533
534
535
	require.False(t, resp)
}

func TestUnloadAllRunners(t *testing.T) {
536
	ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
537
538
539
540
541
542
543
544
545
546
	defer done()

	llm1 := &mockLlm{}
	llm2 := &mockLlm{}
	s := InitScheduler(ctx)
	s.unloadAllRunners()

	r1 := &runnerRef{llama: llm1}
	r2 := &runnerRef{llama: llm2}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
547
	s.loadedMu.Lock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
548
549
	s.loaded["a"] = r1
	s.loaded["b"] = r2
Daniel Hiltgen's avatar
Daniel Hiltgen committed
550
	s.loadedMu.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
551
552
553
554
555
556
557
558
559
	s.unloadAllRunners()

	require.True(t, llm1.closeCalled)
	require.True(t, llm2.closeCalled)
}

func TestUnload(t *testing.T) {
	llm1 := &mockLlm{}
	r1 := &runnerRef{llama: llm1}
560
	r2 := &runnerRef{model: &Model{AdapterPaths: []string{"A"}}}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
561
562
563
	r1.unload()
	require.True(t, llm1.closeCalled)
	r2.unload()
564
	require.Nil(t, r2.model)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
}

type mockLlm struct {
	pingResp          error
	waitResp          error
	completionResp    error
	embeddingResp     []float64
	embeddingRespErr  error
	tokenizeResp      []int
	tokenizeRespErr   error
	detokenizeResp    string
	detonekizeRespErr error
	closeResp         error
	closeCalled       bool
	estimatedVRAM     uint64
580
	estimatedTotal    uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
}

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) Embedding(ctx context.Context, prompt string) ([]float64, error) {
	return s.embeddingResp, 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
}
601
602
func (s *mockLlm) EstimatedVRAM() uint64  { return s.estimatedVRAM }
func (s *mockLlm) EstimatedTotal() uint64 { return s.estimatedTotal }