sched_test.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644
  1. package server
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/binary"
  6. "fmt"
  7. "log/slog"
  8. "os"
  9. "testing"
  10. "time"
  11. "github.com/ollama/ollama/api"
  12. "github.com/ollama/ollama/app/lifecycle"
  13. "github.com/ollama/ollama/envconfig"
  14. "github.com/ollama/ollama/format"
  15. "github.com/ollama/ollama/gpu"
  16. "github.com/ollama/ollama/llm"
  17. "github.com/stretchr/testify/require"
  18. )
  19. func init() {
  20. os.Setenv("OLLAMA_DEBUG", "1")
  21. lifecycle.InitLogging()
  22. }
  23. func TestInitScheduler(t *testing.T) {
  24. ctx, done := context.WithCancel(context.Background())
  25. defer done()
  26. s := InitScheduler(ctx)
  27. s.loadedMu.Lock()
  28. require.NotNil(t, s.loaded)
  29. s.loadedMu.Unlock()
  30. }
  31. func TestLoad(t *testing.T) {
  32. ctx, done := context.WithTimeout(context.Background(), 20*time.Millisecond)
  33. defer done()
  34. s := InitScheduler(ctx)
  35. var ggml *llm.GGML // value not used in tests
  36. req := &LlmRequest{
  37. ctx: ctx,
  38. model: &Model{ModelPath: "foo"},
  39. opts: api.DefaultOptions(),
  40. successCh: make(chan *runnerRef, 1),
  41. errCh: make(chan error, 1),
  42. sessionDuration: 2,
  43. }
  44. // Fail to load model first
  45. s.newServerFn = func(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
  46. return nil, fmt.Errorf("something failed to load model blah")
  47. }
  48. gpus := gpu.GpuInfoList{}
  49. s.load(req, ggml, gpus)
  50. require.Empty(t, req.successCh)
  51. require.Len(t, req.errCh, 1)
  52. s.loadedMu.Lock()
  53. require.Empty(t, s.loaded)
  54. s.loadedMu.Unlock()
  55. err := <-req.errCh
  56. require.Contains(t, err.Error(), "this model may be incompatible")
  57. server := &mockLlm{estimatedVRAM: 10, estimatedVRAMByGPU: map[string]uint64{}}
  58. s.newServerFn = func(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
  59. return server, nil
  60. }
  61. s.load(req, ggml, gpus)
  62. select {
  63. case err := <-req.errCh:
  64. require.NoError(t, err)
  65. case resp := <-req.successCh:
  66. require.Equal(t, uint64(10), resp.estimatedVRAM)
  67. require.Equal(t, uint(1), resp.refCount)
  68. s.loadedMu.Lock()
  69. require.Len(t, s.loaded, 1)
  70. s.loadedMu.Unlock()
  71. }
  72. req.model.ModelPath = "dummy_model_path"
  73. server.waitResp = fmt.Errorf("wait failure")
  74. s.load(req, ggml, gpus)
  75. select {
  76. case err := <-req.errCh:
  77. require.Contains(t, err.Error(), "wait failure")
  78. case resp := <-req.successCh:
  79. t.Errorf("unexpected success %v", resp)
  80. }
  81. s.loadedMu.Lock()
  82. runner := s.loaded["dummy_model_path"]
  83. s.loadedMu.Unlock()
  84. require.NotNil(t, runner)
  85. require.Equal(t, uint(0), runner.refCount)
  86. time.Sleep(1 * time.Millisecond)
  87. require.Len(t, s.expiredCh, 1)
  88. }
  89. type bundle struct {
  90. ctx context.Context //nolint:containedctx
  91. ctxDone func()
  92. srv *mockLlm
  93. req *LlmRequest
  94. ggml *llm.GGML
  95. }
  96. func (scenario *bundle) newServer(gpus gpu.GpuInfoList, model string, ggml *llm.GGML, adapters []string, projectors []string, opts api.Options) (llm.LlamaServer, error) {
  97. return scenario.srv, nil
  98. }
  99. func newScenario(t *testing.T, ctx context.Context, modelName string, estimatedVRAM uint64) *bundle {
  100. scenario := &bundle{}
  101. scenario.ctx, scenario.ctxDone = context.WithCancel(ctx)
  102. t.Helper()
  103. f, err := os.CreateTemp(t.TempDir(), modelName)
  104. require.NoError(t, err)
  105. defer f.Close()
  106. gguf := llm.NewGGUFV3(binary.LittleEndian)
  107. err = gguf.Encode(f, llm.KV{
  108. "general.architecture": "llama",
  109. "general.name": "name",
  110. "llama.context_length": uint32(32),
  111. "llama.embedding_length": uint32(4096),
  112. "llama.block_count": uint32(1),
  113. "llama.attention.head_count": uint32(32),
  114. "llama.attention.head_count_kv": uint32(32),
  115. "tokenizer.ggml.tokens": []string{" "},
  116. "tokenizer.ggml.scores": []float32{0},
  117. "tokenizer.ggml.token_type": []int32{0},
  118. }, []llm.Tensor{
  119. {Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: &bytes.Reader{}},
  120. {Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: &bytes.Reader{}},
  121. })
  122. require.NoError(t, err)
  123. fname := f.Name()
  124. model := &Model{Name: modelName, ModelPath: fname}
  125. scenario.ggml, err = llm.LoadModel(model.ModelPath)
  126. require.NoError(t, err)
  127. scenario.req = &LlmRequest{
  128. ctx: scenario.ctx,
  129. model: model,
  130. opts: api.DefaultOptions(),
  131. sessionDuration: 5 * time.Millisecond,
  132. successCh: make(chan *runnerRef, 1),
  133. errCh: make(chan error, 1),
  134. }
  135. scenario.srv = &mockLlm{estimatedVRAM: estimatedVRAM, estimatedVRAMByGPU: map[string]uint64{"": estimatedVRAM}}
  136. return scenario
  137. }
  138. func TestRequests(t *testing.T) {
  139. ctx, done := context.WithTimeout(context.Background(), time.Second)
  140. defer done()
  141. // Same model, same request
  142. scenario1a := newScenario(t, ctx, "ollama-model-1", 10)
  143. scenario1a.req.sessionDuration = 5 * time.Millisecond
  144. scenario1b := newScenario(t, ctx, "ollama-model-1", 11)
  145. scenario1b.req.model = scenario1a.req.model
  146. scenario1b.ggml = scenario1a.ggml
  147. scenario1b.req.sessionDuration = 0
  148. // simple reload of same model
  149. scenario2a := newScenario(t, ctx, "ollama-model-1", 20)
  150. tmpModel := *scenario1a.req.model
  151. scenario2a.req.model = &tmpModel
  152. scenario2a.ggml = scenario1a.ggml
  153. scenario2a.req.sessionDuration = 5 * time.Millisecond
  154. // Multiple loaded models
  155. scenario3a := newScenario(t, ctx, "ollama-model-3a", 1*format.GigaByte)
  156. scenario3b := newScenario(t, ctx, "ollama-model-3b", 24*format.GigaByte)
  157. scenario3c := newScenario(t, ctx, "ollama-model-4a", 30)
  158. scenario3c.req.opts.NumGPU = 0 // CPU load, will be allowed
  159. scenario3d := newScenario(t, ctx, "ollama-model-3c", 30) // Needs prior unloaded
  160. s := InitScheduler(ctx)
  161. s.getGpuFn = func() gpu.GpuInfoList {
  162. g := gpu.GpuInfo{Library: "metal"}
  163. g.TotalMemory = 24 * format.GigaByte
  164. g.FreeMemory = 12 * format.GigaByte
  165. return []gpu.GpuInfo{g}
  166. }
  167. s.getCpuFn = func() gpu.GpuInfoList {
  168. g := gpu.GpuInfo{Library: "cpu"}
  169. g.TotalMemory = 32 * format.GigaByte
  170. g.FreeMemory = 26 * format.GigaByte
  171. return []gpu.GpuInfo{g}
  172. }
  173. s.newServerFn = scenario1a.newServer
  174. slog.Info("scenario1a")
  175. s.pendingReqCh <- scenario1a.req
  176. require.Len(t, s.pendingReqCh, 1)
  177. s.Run(ctx)
  178. select {
  179. case resp := <-scenario1a.req.successCh:
  180. require.Equal(t, resp.llama, scenario1a.srv)
  181. require.Empty(t, s.pendingReqCh)
  182. require.Empty(t, scenario1a.req.errCh)
  183. case <-ctx.Done():
  184. t.Errorf("timeout")
  185. }
  186. // Same runner as first request due to not needing a reload
  187. s.newServerFn = scenario1b.newServer
  188. slog.Info("scenario1b")
  189. s.pendingReqCh <- scenario1b.req
  190. select {
  191. case resp := <-scenario1b.req.successCh:
  192. require.Equal(t, resp.llama, scenario1a.srv)
  193. require.Empty(t, s.pendingReqCh)
  194. require.Empty(t, scenario1b.req.errCh)
  195. case <-ctx.Done():
  196. t.Errorf("timeout")
  197. }
  198. // Trigger a reload
  199. s.newServerFn = scenario2a.newServer
  200. scenario2a.req.model.AdapterPaths = []string{"new"}
  201. slog.Info("scenario2a")
  202. s.pendingReqCh <- scenario2a.req
  203. // finish first two requests, so model can reload
  204. time.Sleep(1 * time.Millisecond)
  205. scenario1a.ctxDone()
  206. scenario1b.ctxDone()
  207. select {
  208. case resp := <-scenario2a.req.successCh:
  209. require.Equal(t, resp.llama, scenario2a.srv)
  210. require.Empty(t, s.pendingReqCh)
  211. require.Empty(t, scenario2a.req.errCh)
  212. case <-ctx.Done():
  213. t.Errorf("timeout")
  214. }
  215. envconfig.MaxRunners = 1
  216. s.newServerFn = scenario3a.newServer
  217. slog.Info("scenario3a")
  218. s.pendingReqCh <- scenario3a.req
  219. // finish prior request, so new model can load
  220. time.Sleep(1 * time.Millisecond)
  221. scenario2a.ctxDone()
  222. select {
  223. case resp := <-scenario3a.req.successCh:
  224. require.Equal(t, resp.llama, scenario3a.srv)
  225. require.Empty(t, s.pendingReqCh)
  226. require.Empty(t, scenario3a.req.errCh)
  227. case <-ctx.Done():
  228. t.Errorf("timeout")
  229. }
  230. s.loadedMu.Lock()
  231. require.Len(t, s.loaded, 1)
  232. s.loadedMu.Unlock()
  233. envconfig.MaxRunners = 0
  234. s.newServerFn = scenario3b.newServer
  235. slog.Info("scenario3b")
  236. s.pendingReqCh <- scenario3b.req
  237. select {
  238. case resp := <-scenario3b.req.successCh:
  239. require.Equal(t, resp.llama, scenario3b.srv)
  240. require.Empty(t, s.pendingReqCh)
  241. require.Empty(t, scenario3b.req.errCh)
  242. case <-ctx.Done():
  243. t.Errorf("timeout")
  244. }
  245. s.loadedMu.Lock()
  246. require.Len(t, s.loaded, 2)
  247. s.loadedMu.Unlock()
  248. // This is a CPU load with NumGPU = 0 so it should load
  249. s.newServerFn = scenario3c.newServer
  250. slog.Info("scenario3c")
  251. s.pendingReqCh <- scenario3c.req
  252. select {
  253. case resp := <-scenario3c.req.successCh:
  254. require.Equal(t, resp.llama, scenario3c.srv)
  255. require.Empty(t, s.pendingReqCh)
  256. require.Empty(t, scenario3c.req.errCh)
  257. case <-ctx.Done():
  258. t.Errorf("timeout")
  259. }
  260. s.loadedMu.Lock()
  261. require.Len(t, s.loaded, 3)
  262. s.loadedMu.Unlock()
  263. // Try to load a model that wont fit
  264. s.newServerFn = scenario3d.newServer
  265. slog.Info("scenario3d")
  266. s.loadedMu.Lock()
  267. require.Len(t, s.loaded, 3)
  268. s.loadedMu.Unlock()
  269. scenario3a.ctxDone() // Won't help since this one isn't big enough to make room
  270. time.Sleep(2 * time.Millisecond)
  271. s.pendingReqCh <- scenario3d.req
  272. // finish prior request, so new model can load
  273. time.Sleep(6 * time.Millisecond)
  274. s.loadedMu.Lock()
  275. require.Len(t, s.loaded, 2)
  276. s.loadedMu.Unlock()
  277. scenario3b.ctxDone()
  278. select {
  279. case resp := <-scenario3d.req.successCh:
  280. require.Equal(t, resp.llama, scenario3d.srv)
  281. require.Empty(t, s.pendingReqCh)
  282. require.Empty(t, scenario3d.req.errCh)
  283. case <-ctx.Done():
  284. t.Errorf("timeout")
  285. }
  286. s.loadedMu.Lock()
  287. require.Len(t, s.loaded, 2)
  288. s.loadedMu.Unlock()
  289. }
  290. func TestGetRunner(t *testing.T) {
  291. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  292. defer done()
  293. scenario1a := newScenario(t, ctx, "ollama-model-1a", 10)
  294. scenario1a.req.sessionDuration = 0
  295. scenario1b := newScenario(t, ctx, "ollama-model-1b", 10)
  296. scenario1b.req.sessionDuration = 0
  297. scenario1c := newScenario(t, ctx, "ollama-model-1c", 10)
  298. scenario1c.req.sessionDuration = 0
  299. envconfig.MaxQueuedRequests = 1
  300. s := InitScheduler(ctx)
  301. s.getGpuFn = func() gpu.GpuInfoList {
  302. g := gpu.GpuInfo{Library: "metal"}
  303. g.TotalMemory = 24 * format.GigaByte
  304. g.FreeMemory = 12 * format.GigaByte
  305. return []gpu.GpuInfo{g}
  306. }
  307. s.newServerFn = scenario1a.newServer
  308. slog.Info("scenario1a")
  309. successCh1a, errCh1a := s.GetRunner(scenario1a.ctx, scenario1a.req.model, scenario1a.req.opts, scenario1a.req.sessionDuration)
  310. require.Len(t, s.pendingReqCh, 1)
  311. slog.Info("scenario1b")
  312. successCh1b, errCh1b := s.GetRunner(scenario1b.ctx, scenario1b.req.model, scenario1b.req.opts, scenario1b.req.sessionDuration)
  313. require.Len(t, s.pendingReqCh, 1)
  314. require.Empty(t, successCh1b)
  315. require.Len(t, errCh1b, 1)
  316. err := <-errCh1b
  317. require.Contains(t, err.Error(), "server busy")
  318. s.Run(ctx)
  319. select {
  320. case resp := <-successCh1a:
  321. require.Equal(t, resp.llama, scenario1a.srv)
  322. require.Empty(t, s.pendingReqCh)
  323. require.Empty(t, errCh1a)
  324. case <-ctx.Done():
  325. t.Errorf("timeout")
  326. }
  327. scenario1a.ctxDone()
  328. s.loadedMu.Lock()
  329. require.Len(t, s.loaded, 1)
  330. s.loadedMu.Unlock()
  331. scenario1c.req.model.ModelPath = "bad path"
  332. slog.Info("scenario1c")
  333. successCh1c, errCh1c := s.GetRunner(scenario1c.ctx, scenario1c.req.model, scenario1c.req.opts, scenario1c.req.sessionDuration)
  334. // Starts in pending channel, then should be quickly processsed to return an error
  335. time.Sleep(5 * time.Millisecond)
  336. require.Empty(t, successCh1c)
  337. s.loadedMu.Lock()
  338. require.Empty(t, s.loaded)
  339. s.loadedMu.Unlock()
  340. require.Len(t, errCh1c, 1)
  341. err = <-errCh1c
  342. require.Contains(t, err.Error(), "bad path")
  343. scenario1b.ctxDone()
  344. }
  345. // TODO - add one scenario that triggers the bogus finished event with positive ref count
  346. func TestPrematureExpired(t *testing.T) {
  347. ctx, done := context.WithTimeout(context.Background(), 500*time.Millisecond)
  348. defer done()
  349. // Same model, same request
  350. scenario1a := newScenario(t, ctx, "ollama-model-1a", 10)
  351. s := InitScheduler(ctx)
  352. s.getGpuFn = func() gpu.GpuInfoList {
  353. g := gpu.GpuInfo{Library: "metal"}
  354. g.TotalMemory = 24 * format.GigaByte
  355. g.FreeMemory = 12 * format.GigaByte
  356. return []gpu.GpuInfo{g}
  357. }
  358. s.newServerFn = scenario1a.newServer
  359. successCh1a, errCh1a := s.GetRunner(scenario1a.ctx, scenario1a.req.model, scenario1a.req.opts, scenario1a.req.sessionDuration)
  360. require.Len(t, s.pendingReqCh, 1)
  361. s.Run(ctx)
  362. select {
  363. case resp := <-successCh1a:
  364. require.Equal(t, resp.llama, scenario1a.srv)
  365. require.Empty(t, s.pendingReqCh)
  366. require.Empty(t, errCh1a)
  367. s.loadedMu.Lock()
  368. require.Len(t, s.loaded, 1)
  369. s.loadedMu.Unlock()
  370. slog.Info("sending premature expired event now")
  371. s.expiredCh <- resp // Shouldn't happen in real life, but make sure its safe
  372. case <-ctx.Done():
  373. t.Errorf("timeout")
  374. }
  375. time.Sleep(scenario1a.req.sessionDuration)
  376. scenario1a.ctxDone()
  377. time.Sleep(20 * time.Millisecond)
  378. require.LessOrEqual(t, len(s.finishedReqCh), 1)
  379. time.Sleep(10 * time.Millisecond)
  380. require.Empty(t, s.finishedReqCh)
  381. s.loadedMu.Lock()
  382. require.Empty(t, s.loaded)
  383. s.loadedMu.Unlock()
  384. // also shouldn't happen in real life
  385. s.finishedReqCh <- scenario1a.req
  386. time.Sleep(5 * time.Millisecond)
  387. }
  388. func TestUseLoadedRunner(t *testing.T) {
  389. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  390. req := &LlmRequest{
  391. ctx: ctx,
  392. opts: api.DefaultOptions(),
  393. successCh: make(chan *runnerRef, 1),
  394. sessionDuration: 2,
  395. }
  396. finished := make(chan *LlmRequest)
  397. llm1 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{}}
  398. r1 := &runnerRef{llama: llm1, sessionDuration: 1}
  399. req.useLoadedRunner(r1, finished)
  400. require.Equal(t, uint(1), r1.refCount)
  401. require.Equal(t, time.Duration(2), r1.sessionDuration)
  402. select {
  403. case success := <-req.successCh:
  404. require.Equal(t, r1, success)
  405. case <-ctx.Done():
  406. t.Errorf("timeout")
  407. }
  408. done()
  409. fin := <-finished
  410. require.Equal(t, req, fin)
  411. }
  412. func TestUpdateFreeSpace(t *testing.T) {
  413. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  414. defer done()
  415. gpus := gpu.GpuInfoList{
  416. {
  417. Library: "a",
  418. ID: "1",
  419. },
  420. {
  421. Library: "a",
  422. ID: "2",
  423. },
  424. }
  425. gpus[0].TotalMemory = 1000
  426. gpus[0].FreeMemory = 900
  427. gpus[1].TotalMemory = 2000
  428. gpus[1].FreeMemory = 1900
  429. llm1 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{"1": 50, "2": 50}}
  430. llm2 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{"1": 125, "2": 75}}
  431. r1 := &runnerRef{llama: llm1, gpus: gpus}
  432. r2 := &runnerRef{llama: llm2, gpus: gpus}
  433. s := InitScheduler(ctx)
  434. s.loadedMu.Lock()
  435. s.loaded["a"] = r1
  436. s.loaded["b"] = r2
  437. s.loadedMu.Unlock()
  438. s.updateFreeSpace(gpus)
  439. require.Equal(t, uint64(1000-50-125), gpus[0].FreeMemory)
  440. require.Equal(t, uint64(2000-50-75), gpus[1].FreeMemory)
  441. }
  442. func TestFilterGPUsWithoutLoadingModels(t *testing.T) {
  443. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  444. defer done()
  445. gpus := gpu.GpuInfoList{
  446. {
  447. Library: "cuda",
  448. ID: "0",
  449. },
  450. {
  451. Library: "cuda",
  452. ID: "1",
  453. },
  454. }
  455. r1 := &runnerRef{gpus: gpu.GpuInfoList{gpus[0]}, loading: true}
  456. s := InitScheduler(ctx)
  457. s.loadedMu.Lock()
  458. s.loaded["a"] = r1
  459. s.loadedMu.Unlock()
  460. tmp := s.filterGPUsWithoutLoadingModels(gpus)
  461. require.Len(t, tmp, 1)
  462. require.Equal(t, "1", tmp[0].ID)
  463. r1.gpus = gpu.GpuInfoList{gpus[1]}
  464. tmp = s.filterGPUsWithoutLoadingModels(gpus)
  465. require.Len(t, tmp, 1)
  466. require.Equal(t, "0", tmp[0].ID)
  467. r1.gpus = gpu.GpuInfoList{}
  468. tmp = s.filterGPUsWithoutLoadingModels(gpus)
  469. require.Len(t, tmp, 2)
  470. }
  471. func TestFindRunnerToUnload(t *testing.T) {
  472. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  473. defer done()
  474. r1 := &runnerRef{refCount: 1, sessionDuration: 1}
  475. r2 := &runnerRef{sessionDuration: 2}
  476. s := InitScheduler(ctx)
  477. s.loadedMu.Lock()
  478. s.loaded["a"] = r1
  479. s.loaded["b"] = r2
  480. s.loadedMu.Unlock()
  481. resp := s.findRunnerToUnload()
  482. require.Equal(t, r2, resp)
  483. r2.refCount = 1
  484. resp = s.findRunnerToUnload()
  485. require.Equal(t, r1, resp)
  486. }
  487. func TestNeedsReload(t *testing.T) {
  488. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  489. defer done()
  490. llm := &mockLlm{estimatedVRAMByGPU: map[string]uint64{}}
  491. do := api.DefaultOptions()
  492. runner := &runnerRef{
  493. model: &Model{AdapterPaths: []string{"adapter1"}, ProjectorPaths: []string{"projector1"}},
  494. Options: &do,
  495. llama: llm,
  496. }
  497. req := &LlmRequest{
  498. model: &Model{
  499. AdapterPaths: []string{"adapter2"},
  500. ProjectorPaths: []string{"projector2"},
  501. },
  502. opts: api.DefaultOptions(),
  503. }
  504. resp := runner.needsReload(ctx, req)
  505. require.True(t, resp)
  506. req.model.AdapterPaths = runner.model.AdapterPaths
  507. resp = runner.needsReload(ctx, req)
  508. require.True(t, resp)
  509. req.model.ProjectorPaths = runner.model.ProjectorPaths
  510. runner.loading = true
  511. req.opts.NumBatch = 1234
  512. resp = runner.needsReload(ctx, req)
  513. require.True(t, resp)
  514. req.opts.NumBatch = runner.Options.NumBatch
  515. llm.pingResp = fmt.Errorf("foo")
  516. resp = runner.needsReload(ctx, req)
  517. require.True(t, resp)
  518. llm.pingResp = nil
  519. resp = runner.needsReload(ctx, req)
  520. require.False(t, resp)
  521. req.opts.NumGPU = 99
  522. resp = runner.needsReload(ctx, req)
  523. require.True(t, resp)
  524. req.opts.NumGPU = -1
  525. resp = runner.needsReload(ctx, req)
  526. require.False(t, resp)
  527. }
  528. func TestUnloadAllRunners(t *testing.T) {
  529. ctx, done := context.WithTimeout(context.Background(), 100*time.Millisecond)
  530. defer done()
  531. llm1 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{}}
  532. llm2 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{}}
  533. s := InitScheduler(ctx)
  534. s.unloadAllRunners()
  535. r1 := &runnerRef{llama: llm1}
  536. r2 := &runnerRef{llama: llm2}
  537. s.loadedMu.Lock()
  538. s.loaded["a"] = r1
  539. s.loaded["b"] = r2
  540. s.loadedMu.Unlock()
  541. s.unloadAllRunners()
  542. require.True(t, llm1.closeCalled)
  543. require.True(t, llm2.closeCalled)
  544. }
  545. func TestUnload(t *testing.T) {
  546. llm1 := &mockLlm{estimatedVRAMByGPU: map[string]uint64{}}
  547. r1 := &runnerRef{llama: llm1}
  548. r2 := &runnerRef{model: &Model{AdapterPaths: []string{"A"}}}
  549. r1.unload()
  550. require.True(t, llm1.closeCalled)
  551. r2.unload()
  552. require.Nil(t, r2.model)
  553. }
  554. type mockLlm struct {
  555. pingResp error
  556. waitResp error
  557. completionResp error
  558. embeddingResp []float64
  559. embeddingRespErr error
  560. tokenizeResp []int
  561. tokenizeRespErr error
  562. detokenizeResp string
  563. detonekizeRespErr error
  564. closeResp error
  565. closeCalled bool
  566. estimatedVRAM uint64
  567. estimatedTotal uint64
  568. estimatedVRAMByGPU map[string]uint64
  569. }
  570. func (s *mockLlm) Ping(ctx context.Context) error { return s.pingResp }
  571. func (s *mockLlm) WaitUntilRunning(ctx context.Context) error { return s.waitResp }
  572. func (s *mockLlm) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
  573. return s.completionResp
  574. }
  575. func (s *mockLlm) Embedding(ctx context.Context, prompt string) ([]float64, error) {
  576. return s.embeddingResp, s.embeddingRespErr
  577. }
  578. func (s *mockLlm) Tokenize(ctx context.Context, content string) ([]int, error) {
  579. return s.tokenizeResp, s.tokenizeRespErr
  580. }
  581. func (s *mockLlm) Detokenize(ctx context.Context, tokens []int) (string, error) {
  582. return s.detokenizeResp, s.detonekizeRespErr
  583. }
  584. func (s *mockLlm) Close() error {
  585. s.closeCalled = true
  586. return s.closeResp
  587. }
  588. func (s *mockLlm) EstimatedVRAM() uint64 { return s.estimatedVRAM }
  589. func (s *mockLlm) EstimatedTotal() uint64 { return s.estimatedTotal }
  590. func (s *mockLlm) EstimatedVRAMByGPU(gpuid string) uint64 { return s.estimatedVRAMByGPU[gpuid] }