dyn_ext_server.go 10.9 KB
Newer Older
1
2
3
package llm

/*
4
#cgo CFLAGS: -I${SRCDIR}/ext_server -I${SRCDIR}/llama.cpp -I${SRCDIR}/llama.cpp/common -I${SRCDIR}/llama.cpp/examples/server
5
#cgo CFLAGS: -DNDEBUG -DLLAMA_SERVER_LIBRARY=1 -D_XOPEN_SOURCE=600 -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
6
7
#cgo CFLAGS: -Wmissing-noreturn -Wextra -Wcast-qual -Wno-unused-function -Wno-array-bounds
#cgo CPPFLAGS: -Ofast -Wextra -Wno-unused-function -Wno-unused-variable -Wno-deprecated-declarations -Wno-unused-but-set-variable
8
9
#cgo darwin CFLAGS: -D_DARWIN_C_SOURCE
#cgo darwin CPPFLAGS:  -DGGML_USE_ACCELERATE
Daniel Hiltgen's avatar
Daniel Hiltgen committed
10
#cgo darwin CPPFLAGS: -DGGML_USE_METAL -DGGML_METAL_NDEBUG
11
#cgo darwin LDFLAGS: -lc++ -framework Accelerate
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
#cgo darwin LDFLAGS: -framework Foundation -framework Metal -framework MetalKit -framework MetalPerformanceShaders
13
#cgo linux CFLAGS: -D_GNU_SOURCE
14
15
#cgo linux LDFLAGS: -lrt -ldl -lstdc++ -lm
#cgo linux windows LDFLAGS: -lpthread
16
17

#include <stdlib.h>
18
#include "dyn_ext_server.h"
19
20
21

*/
import "C"
22

23
24
25
26
27
import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
28
	"log/slog"
29
30
31
	"os"
	"path/filepath"
	"runtime"
32
	"strings"
33
34
35
36
37
38
39
	"sync"
	"time"
	"unsafe"

	"github.com/jmorganca/ollama/api"
)

40
41
42
type dynExtServer struct {
	s       C.struct_dynamic_llama_server
	options api.Options
43
44
45
46
47
}

// Note: current implementation does not support concurrent instantiations
var mutex sync.Mutex

48
49
50
51
52
53
func newExtServerResp(len C.size_t) C.ext_server_resp_t {
	var resp C.ext_server_resp_t
	resp.msg_len = len
	bytes := make([]byte, len)
	resp.msg = (*C.char)(C.CBytes(bytes))
	return resp
54
55
}

56
57
58
59
60
func freeExtServerResp(resp C.ext_server_resp_t) {
	if resp.msg_len == 0 {
		return
	}
	C.free(unsafe.Pointer(resp.msg))
61
62
}

63
64
func extServerResponseToErr(resp C.ext_server_resp_t) error {
	return fmt.Errorf(C.GoString(resp.msg))
65
66
}

67
68
69
70
// Note: current implementation does not support concurrent instantiations
var llm *dynExtServer

func newDynExtServer(library, model string, adapters, projectors []string, opts api.Options) (LLM, error) {
71
	if !mutex.TryLock() {
72
		slog.Info("concurrent llm servers not yet supported, waiting for prior server to complete")
73
74
		mutex.Lock()
	}
75
76
77
	updatePath(filepath.Dir(library))
	libPath := C.CString(library)
	defer C.free(unsafe.Pointer(libPath))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
78
	resp := newExtServerResp(512)
79
80
81
82
83
84
85
86
87
88
89
	defer freeExtServerResp(resp)
	var srv C.struct_dynamic_llama_server
	C.dyn_init(libPath, &srv, &resp)
	if resp.id < 0 {
		mutex.Unlock()
		return nil, fmt.Errorf("Unable to load dynamic library: %s", C.GoString(resp.msg))
	}
	llm = &dynExtServer{
		s:       srv,
		options: opts,
	}
90
	slog.Info(fmt.Sprintf("Loading Dynamic llm server: %s", library))
91

92
	var sparams C.ext_server_params_t
93
94
95
96
97
98
	sparams.model = C.CString(model)
	defer C.free(unsafe.Pointer(sparams.model))

	sparams.embedding = true
	sparams.n_ctx = C.uint(opts.NumCtx)
	sparams.n_batch = C.uint(opts.NumBatch)
99
	sparams.n_gpu_layers = C.int(opts.NumGPU)
100
	sparams.main_gpu = C.int(opts.MainGPU)
101
	sparams.n_parallel = 1 // TODO - wire up concurrency
102
103
104
105

	// Always use the value encoded in the model
	sparams.rope_freq_base = 0.0
	sparams.rope_freq_scale = 0.0
106
107
108
109
	sparams.memory_f16 = C.bool(opts.F16KV)
	sparams.use_mlock = C.bool(opts.UseMLock)
	sparams.use_mmap = C.bool(opts.UseMMap)
	sparams.numa = C.bool(opts.UseNUMA)
110
111
112

	sparams.lora_adapters = nil
	for i := 0; i < len(adapters); i++ {
113
		la := (*C.ext_server_lora_adapter_t)(C.malloc(C.sizeof_ext_server_lora_adapter_t))
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
		defer C.free(unsafe.Pointer(la))
		la.adapter = C.CString(adapters[i])
		defer C.free(unsafe.Pointer(la.adapter))
		la.scale = C.float(1.0) // TODO expose scale/weights up through ollama UX
		la.next = nil
		if i == 0 {
			sparams.lora_adapters = la
		} else {
			tmp := sparams.lora_adapters
			for ; tmp.next != nil; tmp = tmp.next {
			}
			tmp.next = la
		}
	}

129
130
131
132
133
134
135
	if len(projectors) > 0 {
		// TODO: applying multiple projectors is not supported by the llama.cpp server yet
		sparams.mmproj = C.CString(projectors[0])
		defer C.free(unsafe.Pointer(sparams.mmproj))
	} else {
		sparams.mmproj = nil
	}
136

137
	sparams.n_threads = C.uint(opts.NumThread)
138

139
	slog.Info("Initializing llama server")
140
141
142
143
144
	initResp := newExtServerResp(128)
	defer freeExtServerResp(initResp)
	C.dyn_llama_server_init(llm.s, &sparams, &initResp)
	if initResp.id < 0 {
		return nil, extServerResponseToErr(initResp)
145
146
	}

147
	slog.Info("Starting llama main loop")
148
149
	C.dyn_llama_server_start(llm.s)
	return llm, nil
150
151
}

152
func (llm *dynExtServer) Predict(ctx context.Context, predict PredictOpts, fn func(PredictResult)) error {
153
154
155
156
157
158
159
160
	resp := newExtServerResp(128)
	defer freeExtServerResp(resp)
	var imageData []ImageData
	if len(predict.Images) > 0 {
		for cnt, i := range predict.Images {
			imageData = append(imageData, ImageData{Data: i, ID: cnt})
		}
	}
161
	slog.Info(fmt.Sprintf("loaded %d images", len(imageData)))
162
163
164
165

	request := map[string]any{
		"prompt":            predict.Prompt,
		"stream":            true,
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
		"n_predict":         predict.Options.NumPredict,
		"n_keep":            predict.Options.NumKeep,
		"temperature":       predict.Options.Temperature,
		"top_k":             predict.Options.TopK,
		"top_p":             predict.Options.TopP,
		"tfs_z":             predict.Options.TFSZ,
		"typical_p":         predict.Options.TypicalP,
		"repeat_last_n":     predict.Options.RepeatLastN,
		"repeat_penalty":    predict.Options.RepeatPenalty,
		"presence_penalty":  predict.Options.PresencePenalty,
		"frequency_penalty": predict.Options.FrequencyPenalty,
		"mirostat":          predict.Options.Mirostat,
		"mirostat_tau":      predict.Options.MirostatTau,
		"mirostat_eta":      predict.Options.MirostatEta,
		"penalize_nl":       predict.Options.PenalizeNewline,
		"seed":              predict.Options.Seed,
		"stop":              predict.Options.Stop,
183
		"image_data":        imageData,
184
185
186
187
188
189
	}

	if predict.Format == "json" {
		request["grammar"] = jsonGrammar
	}

190
191
192
193
194
195
	retryDelay := 100 * time.Microsecond
	for retries := 0; retries < maxRetries; retries++ {
		if retries > 0 {
			time.Sleep(retryDelay) // wait before retrying
			retryDelay *= 2        // exponential backoff
		}
196

197
198
199
200
		// Handling JSON marshaling with special characters unescaped.
		buffer := &bytes.Buffer{}
		enc := json.NewEncoder(buffer)
		enc.SetEscapeHTML(false)
201

202
203
204
		if err := enc.Encode(request); err != nil {
			return fmt.Errorf("failed to marshal data: %w", err)
		}
205

206
207
		req := C.CString(buffer.String())
		defer C.free(unsafe.Pointer(req))
208

209
		C.dyn_llama_server_completion(llm.s, req, &resp)
210
211
212
		if resp.id < 0 {
			return extServerResponseToErr(resp)
		}
213

214
215
216
217
218
219
		retryNeeded := false
	out:
		for {
			select {
			case <-ctx.Done():
				// This handles the request cancellation
220
				C.dyn_llama_server_completion_cancel(llm.s, resp.id, &resp)
221
222
223
224
225
226
227
				if resp.id < 0 {
					return extServerResponseToErr(resp)
				} else {
					return nil
				}
			default:
				var result C.ext_server_task_result_t
228
				C.dyn_llama_server_completion_next_result(llm.s, resp.id, &result)
229
				json_resp := C.GoString(result.json_resp)
230
				C.dyn_llama_server_release_task_result(llm.s, &result)
231
232
233

				var p prediction
				if err := json.Unmarshal([]byte(json_resp), &p); err != nil {
234
					C.dyn_llama_server_completion_cancel(llm.s, resp.id, &resp)
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
					if resp.id < 0 {
						return fmt.Errorf("error unmarshaling llm prediction response: %w and cancel %s", err, C.GoString(resp.msg))
					} else {
						return fmt.Errorf("error unmarshaling llm prediction response: %w", err)
					}
				}

				if bool(result.error) && strings.Contains(json_resp, "slot unavailable") {
					retryNeeded = true
					// task will already be canceled
					break out
				}

				if p.Content != "" {
					fn(PredictResult{
						Content: p.Content,
					})
				}

				if p.Stop {
					fn(PredictResult{
						Done:               true,
						PromptEvalCount:    p.Timings.PromptN,
						PromptEvalDuration: parseDurationMs(p.Timings.PromptMS),
						EvalCount:          p.Timings.PredictedN,
						EvalDuration:       parseDurationMs(p.Timings.PredictedMS),
					})
					return nil
				}
264
265
			}
		}
266
267
268
		if !retryNeeded {
			return nil // success
		}
269
270
	}

271
272
273
274
	// should never reach here ideally
	return fmt.Errorf("max retries exceeded")
}

275
func (llm *dynExtServer) Encode(ctx context.Context, prompt string) ([]int, error) {
276
277
278
279
280
281
	data, err := json.Marshal(TokenizeRequest{Content: prompt})
	if err != nil {
		return nil, fmt.Errorf("marshaling encode data: %w", err)
	}
	req := C.CString(string(data))
	defer C.free(unsafe.Pointer(req))
282
283
284
	var json_resp *C.char
	resp := newExtServerResp(128)
	defer freeExtServerResp(resp)
285
	C.dyn_llama_server_tokenize(llm.s, req, &json_resp, &resp)
286
287
	if resp.id < 0 {
		return nil, extServerResponseToErr(resp)
288
	}
289
	defer C.dyn_llama_server_release_json_resp(llm.s, &json_resp)
290
291

	var encoded TokenizeResponse
292
	if err2 := json.Unmarshal([]byte(C.GoString(json_resp)), &encoded); err2 != nil {
293
294
295
296
297
298
		return nil, fmt.Errorf("unmarshal encode response: %w", err2)
	}

	return encoded.Tokens, err
}

299
func (llm *dynExtServer) Decode(ctx context.Context, tokens []int) (string, error) {
300
301
302
303
304
305
306
307
308
309
	if len(tokens) == 0 {
		return "", nil
	}
	data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
	if err != nil {
		return "", fmt.Errorf("marshaling decode data: %w", err)
	}

	req := C.CString(string(data))
	defer C.free(unsafe.Pointer(req))
310
311
312
	var json_resp *C.char
	resp := newExtServerResp(128)
	defer freeExtServerResp(resp)
313
	C.dyn_llama_server_detokenize(llm.s, req, &json_resp, &resp)
314
315
	if resp.id < 0 {
		return "", extServerResponseToErr(resp)
316
	}
317
	defer C.dyn_llama_server_release_json_resp(llm.s, &json_resp)
318
319

	var decoded DetokenizeResponse
320
	if err2 := json.Unmarshal([]byte(C.GoString(json_resp)), &decoded); err2 != nil {
321
322
323
324
325
326
		return "", fmt.Errorf("unmarshal encode response: %w", err2)
	}

	return decoded.Content, err
}

327
func (llm *dynExtServer) Embedding(ctx context.Context, input string) ([]float64, error) {
328
329
330
331
332
333
334
	data, err := json.Marshal(TokenizeRequest{Content: input})
	if err != nil {
		return nil, fmt.Errorf("error marshaling embed data: %w", err)
	}

	req := C.CString(string(data))
	defer C.free(unsafe.Pointer(req))
335
336
337
	var json_resp *C.char
	resp := newExtServerResp(128)
	defer freeExtServerResp(resp)
338
	C.dyn_llama_server_embedding(llm.s, req, &json_resp, &resp)
339
340
	if resp.id < 0 {
		return nil, extServerResponseToErr(resp)
341
	}
342
	defer C.dyn_llama_server_release_json_resp(llm.s, &json_resp)
343
344

	var embedding EmbeddingResponse
345
	if err := json.Unmarshal([]byte(C.GoString(json_resp)), &embedding); err != nil {
346
347
348
349
350
351
		return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
	}

	return embedding.Embedding, nil
}

352
353
func (llm *dynExtServer) Close() {
	C.dyn_llama_server_stop(llm.s)
354
355
	mutex.Unlock()
}
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372

func updatePath(dir string) {
	if runtime.GOOS == "windows" {
		tmpDir := filepath.Dir(dir)
		pathComponents := strings.Split(os.Getenv("PATH"), ";")
		i := 0
		for _, comp := range pathComponents {
			if strings.EqualFold(comp, dir) {
				return
			}
			// Remove any other prior paths to our temp dir
			if !strings.HasPrefix(strings.ToLower(comp), strings.ToLower(tmpDir)) {
				pathComponents[i] = comp
				i++
			}
		}
		newPath := strings.Join(append([]string{dir}, pathComponents...), ";")
373
		slog.Info(fmt.Sprintf("Updating PATH to %s", newPath))
374
375
		os.Setenv("PATH", newPath)
	}
376
	// linux and darwin rely on rpath
377
}