utils_test.go 10.6 KB
Newer Older
1
2
3
4
5
6
7
//go:build integration

package integration

import (
	"bytes"
	"context"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
8
	"errors"
9
10
11
12
13
14
	"fmt"
	"io"
	"log/slog"
	"math/rand"
	"net"
	"net/http"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
15
	"net/url"
16
17
18
19
20
21
22
23
24
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

25
26
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/app/lifecycle"
27
	"github.com/ollama/ollama/format"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
28
	"github.com/stretchr/testify/require"
29
30
)

31
32
33
34
const (
	smol = "llama3.2:1b"
)

35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
var (
	started = time.Now()

	// Note: add newer models at the top of the list to test them first
	ollamaEngineChatModels = []string{
		"gemma3n:e2b",
		"mistral-small3.2:latest",
		"deepseek-r1:1.5b",
		"llama3.2-vision:latest",
		"qwen2.5-coder:latest",
		"qwen2.5vl:3b",
		"qwen3:0.6b", // dense
		"qwen3:30b",  // MOE
		"gemma3:1b",
		"llama3.1:latest",
		"llama3.2:latest",
		"gemma2:latest",
		"minicpm-v:latest",    // arch=qwen2
		"granite-code:latest", // arch=llama
	}
	llamaRunnerChatModels = []string{
		"mistral:latest",
		"falcon3:latest",
		"granite3-moe:latest",
		"command-r:latest",
		"nemotron-mini:latest",
		"phi3.5:latest",
		"solar-pro:latest",
		"internlm2:latest",
		"codellama:latest", // arch=llama
		"phi3:latest",
		"falcon2:latest",
		"gemma:latest",
		"llama2:latest",
		"nous-hermes:latest",
		"orca-mini:latest",
		"qwen:latest",
		"stablelm2:latest", // Predictions are off, crashes on small VRAM GPUs
		"falcon:latest",
	}
)

Daniel Hiltgen's avatar
Daniel Hiltgen committed
77
78
79
80
func Init() {
	lifecycle.InitLogging()
}

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
func FindPort() string {
	port := 0
	if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
		var l *net.TCPListener
		if l, err = net.ListenTCP("tcp", a); err == nil {
			port = l.Addr().(*net.TCPAddr).Port
			l.Close()
		}
	}
	if port == 0 {
		port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
	}
	return strconv.Itoa(port)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
96
func GetTestEndpoint() (*api.Client, string) {
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
	defaultPort := "11434"
	ollamaHost := os.Getenv("OLLAMA_HOST")

	scheme, hostport, ok := strings.Cut(ollamaHost, "://")
	if !ok {
		scheme, hostport = "http", ollamaHost
	}

	// trim trailing slashes
	hostport = strings.TrimRight(hostport, "/")

	host, port, err := net.SplitHostPort(hostport)
	if err != nil {
		host, port = "127.0.0.1", defaultPort
		if ip := net.ParseIP(strings.Trim(hostport, "[]")); ip != nil {
			host = ip.String()
		} else if hostport != "" {
			host = hostport
		}
	}

	if os.Getenv("OLLAMA_TEST_EXISTING") == "" && port == defaultPort {
		port = FindPort()
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
122
123
124
125
126
127
128
129
	slog.Info("server connection", "host", host, "port", port)

	return api.NewClient(
		&url.URL{
			Scheme: scheme,
			Host:   net.JoinHostPort(host, port),
		},
		http.DefaultClient), fmt.Sprintf("%s:%s", host, port)
130
131
132
133
134
}

var serverMutex sync.Mutex
var serverReady bool

Daniel Hiltgen's avatar
Daniel Hiltgen committed
135
func startServer(t *testing.T, ctx context.Context, ollamaHost string) error {
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
	// Make sure the server has been built
	CLIName, err := filepath.Abs("../ollama")
	if err != nil {
		return err
	}

	if runtime.GOOS == "windows" {
		CLIName += ".exe"
	}
	_, err = os.Stat(CLIName)
	if err != nil {
		return fmt.Errorf("CLI missing, did you forget to build first?  %w", err)
	}
	serverMutex.Lock()
	defer serverMutex.Unlock()
	if serverReady {
		return nil
	}

	if tmp := os.Getenv("OLLAMA_HOST"); tmp != ollamaHost {
		slog.Info("setting env", "OLLAMA_HOST", ollamaHost)
Michael Yang's avatar
Michael Yang committed
157
		t.Setenv("OLLAMA_HOST", ollamaHost)
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
	}

	slog.Info("starting server", "url", ollamaHost)
	done, err := lifecycle.SpawnServer(ctx, "../ollama")
	if err != nil {
		return fmt.Errorf("failed to start server: %w", err)
	}

	go func() {
		<-ctx.Done()
		serverMutex.Lock()
		defer serverMutex.Unlock()
		exitCode := <-done
		if exitCode > 0 {
			slog.Warn("server failure", "exit", exitCode)
		}
		serverReady = false
	}()

	// TODO wait only long enough for the server to be responsive...
	time.Sleep(500 * time.Millisecond)

	serverReady = true
	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
184
func PullIfMissing(ctx context.Context, client *api.Client, modelName string) error {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
185
	slog.Info("checking status of model", "model", modelName)
186
187
	showReq := &api.ShowRequest{Name: modelName}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
188
189
	showCtx, cancel := context.WithDeadlineCause(
		ctx,
190
		time.Now().Add(20*time.Second),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
191
192
193
194
195
196
197
198
199
		fmt.Errorf("show for existing model %s took too long", modelName),
	)
	defer cancel()
	_, err := client.Show(showCtx, showReq)
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
		break
	case err != nil:
200
		return err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
201
	default:
202
203
204
		slog.Info("model already present", "model", modelName)
		return nil
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
205
206
	slog.Info("model missing", "model", modelName)

207
	stallDuration := 60 * time.Second // This includes checksum verification, which can take a while on larger models, and slower systems
Daniel Hiltgen's avatar
Daniel Hiltgen committed
208
209
210
211
	stallTimer := time.NewTimer(stallDuration)
	fn := func(resp api.ProgressResponse) error {
		// fmt.Print(".")
		if !stallTimer.Reset(stallDuration) {
Michael Yang's avatar
lint  
Michael Yang committed
212
			return errors.New("stall was detected, aborting status reporting")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
213
214
215
		}
		return nil
	}
216

Daniel Hiltgen's avatar
Daniel Hiltgen committed
217
	stream := true
218
219
	pullReq := &api.PullRequest{Name: modelName, Stream: &stream}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
220
	var pullError error
221

Daniel Hiltgen's avatar
Daniel Hiltgen committed
222
223
224
225
226
227
228
229
	done := make(chan int)
	go func() {
		pullError = client.Pull(ctx, pullReq, fn)
		done <- 0
	}()

	select {
	case <-stallTimer.C:
Michael Yang's avatar
lint  
Michael Yang committed
230
		return errors.New("download stalled")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
231
232
	case <-done:
		return pullError
233
234
235
	}
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
236
237
var serverProcMutex sync.Mutex

Daniel Hiltgen's avatar
Daniel Hiltgen committed
238
239
240
241
242
243
244
245
246
247
248
249
// Returns an Client, the testEndpoint, and a cleanup function, fails the test on errors
// Starts the server if needed
func InitServerConnection(ctx context.Context, t *testing.T) (*api.Client, string, func()) {
	client, testEndpoint := GetTestEndpoint()
	if os.Getenv("OLLAMA_TEST_EXISTING") == "" {
		serverProcMutex.Lock()
		fp, err := os.CreateTemp("", "ollama-server-*.log")
		if err != nil {
			t.Fatalf("failed to generate log file: %s", err)
		}
		lifecycle.ServerLogFile = fp.Name()
		fp.Close()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
250
		require.NoError(t, startServer(t, ctx, testEndpoint))
251
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
252
253

	return client, testEndpoint, func() {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
254
255
256
257
258
259
260
261
		if os.Getenv("OLLAMA_TEST_EXISTING") == "" {
			defer serverProcMutex.Unlock()
			if t.Failed() {
				fp, err := os.Open(lifecycle.ServerLogFile)
				if err != nil {
					slog.Error("failed to open server log", "logfile", lifecycle.ServerLogFile, "error", err)
					return
				}
262
				defer fp.Close()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
263
264
265
266
267
268
269
270
				data, err := io.ReadAll(fp)
				if err != nil {
					slog.Error("failed to read server log", "logfile", lifecycle.ServerLogFile, "error", err)
					return
				}
				slog.Warn("SERVER LOG FOLLOWS")
				os.Stderr.Write(data)
				slog.Warn("END OF SERVER")
271
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
272
			err := os.Remove(lifecycle.ServerLogFile)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
273
274
			if err != nil && !os.IsNotExist(err) {
				slog.Warn("failed to cleanup", "logfile", lifecycle.ServerLogFile, "error", err)
275
276
277
			}
		}
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
278
}
279

Daniel Hiltgen's avatar
Daniel Hiltgen committed
280
281
282
283
284
285
func GenerateTestHelper(ctx context.Context, t *testing.T, genReq api.GenerateRequest, anyResp []string) {
	client, _, cleanup := InitServerConnection(ctx, t)
	defer cleanup()
	require.NoError(t, PullIfMissing(ctx, client, genReq.Model))
	DoGenerate(ctx, t, client, genReq, anyResp, 30*time.Second, 10*time.Second)
}
286

Daniel Hiltgen's avatar
Daniel Hiltgen committed
287
288
289
290
291
292
293
func DoGenerate(ctx context.Context, t *testing.T, client *api.Client, genReq api.GenerateRequest, anyResp []string, initialTimeout, streamTimeout time.Duration) {
	stallTimer := time.NewTimer(initialTimeout)
	var buf bytes.Buffer
	fn := func(response api.GenerateResponse) error {
		// fmt.Print(".")
		buf.Write([]byte(response.Response))
		if !stallTimer.Reset(streamTimeout) {
Michael Yang's avatar
lint  
Michael Yang committed
294
			return errors.New("stall was detected while streaming response, aborting")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
295
296
		}
		return nil
297
298
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
299
300
301
302
303
304
305
306
	stream := true
	genReq.Stream = &stream
	done := make(chan int)
	var genErr error
	go func() {
		genErr = client.Generate(ctx, &genReq, fn)
		done <- 0
	}()
307

Daniel Hiltgen's avatar
Daniel Hiltgen committed
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
	select {
	case <-stallTimer.C:
		if buf.Len() == 0 {
			t.Errorf("generate never started.  Timed out after :%s", initialTimeout.String())
		} else {
			t.Errorf("generate stalled.  Response so far:%s", buf.String())
		}
	case <-done:
		require.NoError(t, genErr, "failed with %s request prompt %s ", genReq.Model, genReq.Prompt)
		// Verify the response contains the expected data
		response := buf.String()
		atLeastOne := false
		for _, resp := range anyResp {
			if strings.Contains(strings.ToLower(response), resp) {
				atLeastOne = true
				break
			}
		}
326
		require.True(t, atLeastOne, "%s: none of %v found in %s", genReq.Model, anyResp, response)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
327
328
329
		slog.Info("test pass", "model", genReq.Model, "prompt", genReq.Prompt, "contains", anyResp, "response", response)
	case <-ctx.Done():
		t.Error("outer test context done while waiting for generate")
330
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
331
}
332

Daniel Hiltgen's avatar
Daniel Hiltgen committed
333
// Generate a set of requests
334
// By default each request uses llama3.2 as the model
Daniel Hiltgen's avatar
Daniel Hiltgen committed
335
336
337
func GenerateRequests() ([]api.GenerateRequest, [][]string) {
	return []api.GenerateRequest{
			{
338
				Model:     smol,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
339
340
341
				Prompt:    "why is the ocean blue?",
				Stream:    &stream,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
342
				Options: map[string]any{
Daniel Hiltgen's avatar
Daniel Hiltgen committed
343
344
345
346
					"seed":        42,
					"temperature": 0.0,
				},
			}, {
347
				Model:     smol,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
348
349
350
				Prompt:    "why is the color of dirt brown?",
				Stream:    &stream,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
351
				Options: map[string]any{
Daniel Hiltgen's avatar
Daniel Hiltgen committed
352
353
354
355
					"seed":        42,
					"temperature": 0.0,
				},
			}, {
356
				Model:     smol,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
357
358
359
				Prompt:    "what is the origin of the us thanksgiving holiday?",
				Stream:    &stream,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
360
				Options: map[string]any{
Daniel Hiltgen's avatar
Daniel Hiltgen committed
361
362
363
364
					"seed":        42,
					"temperature": 0.0,
				},
			}, {
365
				Model:     smol,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
366
367
368
				Prompt:    "what is the origin of independence day?",
				Stream:    &stream,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
369
				Options: map[string]any{
Daniel Hiltgen's avatar
Daniel Hiltgen committed
370
371
372
373
					"seed":        42,
					"temperature": 0.0,
				},
			}, {
374
				Model:     smol,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
375
376
377
				Prompt:    "what is the composition of air?",
				Stream:    &stream,
				KeepAlive: &api.Duration{Duration: 10 * time.Second},
378
				Options: map[string]any{
Daniel Hiltgen's avatar
Daniel Hiltgen committed
379
380
381
382
383
384
					"seed":        42,
					"temperature": 0.0,
				},
			},
		},
		[][]string{
Michael Yang's avatar
Michael Yang committed
385
386
387
388
389
			{"sunlight"},
			{"soil", "organic", "earth", "black", "tan"},
			{"england", "english", "massachusetts", "pilgrims", "british"},
			{"fourth", "july", "declaration", "independence"},
			{"nitrogen", "oxygen", "carbon", "dioxide"},
390
		}
391
}
392
393
394
395
396
397
398
399
400
401
402
403

func skipUnderMinVRAM(t *testing.T, gb uint64) {
	// TODO use info API in the future
	if s := os.Getenv("OLLAMA_MAX_VRAM"); s != "" {
		maxVram, err := strconv.ParseUint(s, 10, 64)
		require.NoError(t, err)
		// Don't hammer on small VRAM cards...
		if maxVram < gb*format.GibiByte {
			t.Skip("skipping with small VRAM to avoid timeouts")
		}
	}
}
404
405
406
407
408
409
410
411
412
413
414

func getTimeouts(t *testing.T) (soft time.Duration, hard time.Duration) {
	deadline, hasDeadline := t.Deadline()
	if !hasDeadline {
		return 8 * time.Minute, 10 * time.Minute
	} else if deadline.Compare(time.Now().Add(2*time.Minute)) <= 0 {
		t.Skip("too little time")
		return time.Duration(0), time.Duration(0)
	}
	return -time.Since(deadline.Add(-2 * time.Minute)), -time.Since(deadline.Add(-20 * time.Second))
}