shim_ext_server.go 5.6 KB
Newer Older
1
2
3
4
5
6
7
//go:build !darwin

package llm

/*

#include <stdlib.h>
8
#include "dynamic_shim.h"
9
10
11
12
13
14
15

*/
import "C"
import (
	"context"
	"errors"
	"fmt"
16
	"io"
17
18
19
20
	"io/fs"
	"log"
	"os"
	"path/filepath"
21
	"strings"
22
23
24
25
26
27
28
	"sync"
	"unsafe"

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

type shimExtServer struct {
29
	s       C.struct_dynamic_llama_server
30
31
32
33
34
35
36
37
	options api.Options
}

// Note: current implementation does not support concurrent instantiations
var shimMutex sync.Mutex
var llm *shimExtServer

func (llm *shimExtServer) llama_server_init(sparams *C.ext_server_params_t, err *C.ext_server_resp_t) {
38
	C.dynamic_shim_llama_server_init(llm.s, sparams, err)
39
40
}
func (llm *shimExtServer) llama_server_start() {
41
	C.dynamic_shim_llama_server_start(llm.s)
42
43
}
func (llm *shimExtServer) llama_server_stop() {
44
	C.dynamic_shim_llama_server_stop(llm.s)
45
46
47
}

func (llm *shimExtServer) llama_server_completion(json_req *C.char, resp *C.ext_server_resp_t) {
48
	C.dynamic_shim_llama_server_completion(llm.s, json_req, resp)
49
50
}
func (llm *shimExtServer) llama_server_completion_next_result(task_id C.int, resp *C.ext_server_task_result_t) {
51
	C.dynamic_shim_llama_server_completion_next_result(llm.s, task_id, resp)
52
53
}
func (llm *shimExtServer) llama_server_completion_cancel(task_id C.int, err *C.ext_server_resp_t) {
54
	C.dynamic_shim_llama_server_completion_cancel(llm.s, task_id, err)
55
56
}
func (llm *shimExtServer) llama_server_release_task_result(result *C.ext_server_task_result_t) {
57
	C.dynamic_shim_llama_server_release_task_result(llm.s, result)
58
59
60
}

func (llm *shimExtServer) llama_server_tokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
61
	C.dynamic_shim_llama_server_tokenize(llm.s, json_req, json_resp, err)
62
63
}
func (llm *shimExtServer) llama_server_detokenize(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
64
	C.dynamic_shim_llama_server_detokenize(llm.s, json_req, json_resp, err)
65
66
}
func (llm *shimExtServer) llama_server_embedding(json_req *C.char, json_resp **C.char, err *C.ext_server_resp_t) {
67
	C.dynamic_shim_llama_server_embedding(llm.s, json_req, json_resp, err)
68
69
}
func (llm *shimExtServer) llama_server_release_json_resp(json_resp **C.char) {
70
	C.dynamic_shim_llama_server_release_json_resp(llm.s, json_resp)
71
72
}

73
74
75
func newDynamicShimExtServer(library, model string, adapters, projectors []string, numLayers int64, opts api.Options) (extServer, error) {
	shimMutex.Lock()
	defer shimMutex.Unlock()
76
	updatePath(filepath.Dir(library))
77
78
79
80
81
82
83
84
	libPath := C.CString(library)
	defer C.free(unsafe.Pointer(libPath))
	resp := newExtServerResp(128)
	defer freeExtServerResp(resp)
	var srv C.struct_dynamic_llama_server
	C.dynamic_shim_init(libPath, &srv, &resp)
	if resp.id < 0 {
		return nil, fmt.Errorf("Unable to load dynamic library: %s", C.GoString(resp.msg))
85
	}
86
87
88
	llm = &shimExtServer{
		s:       srv,
		options: opts,
89
	}
90
	log.Printf("Loading Dynamic Shim llm server: %s", library)
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
	return newExtServer(llm, model, adapters, projectors, numLayers, opts)
}

func (llm *shimExtServer) Predict(ctx context.Context, pred PredictOpts, fn func(PredictResult)) error {
	return predict(llm, llm.options, ctx, pred, fn)
}

func (llm *shimExtServer) Encode(ctx context.Context, prompt string) ([]int, error) {
	return encode(llm, ctx, prompt)
}

func (llm *shimExtServer) Decode(ctx context.Context, tokens []int) (string, error) {
	return decode(llm, ctx, tokens)
}

func (llm *shimExtServer) Embedding(ctx context.Context, input string) ([]float64, error) {
	return embedding(llm, ctx, input)
}

func (llm *shimExtServer) Close() {
	close(llm)
}

func nativeInit(workdir string) error {
115
	libs, err := extractDynamicLibs(workdir, "llama.cpp/gguf/build/*/*/lib/*")
116
	if err != nil {
117
		if err == payloadMissing {
118
			log.Printf("%s", payloadMissing)
119
120
121
			return nil
		}
		return err
122
123
	}
	for _, lib := range libs {
124
125
126
127
128
129
130
		// The last dir component is the variant name
		variant := filepath.Base(filepath.Dir(lib))
		AvailableShims[variant] = lib
	}

	if err := verifyDriverAccess(); err != nil {
		return err
131
132
	}

133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
	// Report which dynamic libraries we have loaded to assist troubleshooting
	variants := make([]string, len(AvailableShims))
	i := 0
	for variant := range AvailableShims {
		variants[i] = variant
		i++
	}
	log.Printf("Dynamic LLM variants %v", variants)

	return nil
}

func extractDynamicLibs(workDir, glob string) ([]string, error) {
	files, err := fs.Glob(libEmbed, glob)
	if err != nil || len(files) == 0 {
		return nil, payloadMissing
	}
	libs := make([]string, len(files))

	for i, file := range files {
		pathComps := strings.Split(file, "/")
		if len(pathComps) != 7 {
			log.Printf("unexpected payload components: %v", pathComps)
			continue
		}
		// llama.cpp/gguf/build/$OS/$VARIANT/lib/$LIBRARY
		// Include the variant in the path to avoid conflicts between multiple server libs
		targetDir := filepath.Join(workDir, pathComps[4])
		srcFile, err := libEmbed.Open(file)
162
		if err != nil {
163
164
165
166
167
168
			return nil, fmt.Errorf("read payload %s: %v", file, err)
		}
		defer srcFile.Close()
		if err := os.MkdirAll(targetDir, 0o755); err != nil {
			return nil, fmt.Errorf("create payload temp dir %s: %v", workDir, err)
		}
169

170
171
172
		destFile := filepath.Join(targetDir, filepath.Base(file))
		if strings.Contains(destFile, "server") {
			libs[i] = destFile
173
174
		}

175
176
177
178
179
180
181
182
183
184
185
186
187
188
		_, err = os.Stat(destFile)
		switch {
		case errors.Is(err, os.ErrNotExist):
			destFile, err := os.OpenFile(destFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
			if err != nil {
				return nil, fmt.Errorf("write payload %s: %v", file, err)
			}
			defer destFile.Close()
			if _, err := io.Copy(destFile, srcFile); err != nil {
				return nil, fmt.Errorf("copy payload %s: %v", file, err)
			}
		case err != nil:
			return nil, fmt.Errorf("stat payload %s: %v", file, err)
		}
189
	}
190
	return libs, nil
191
}