char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

unraidofix: rpc-server cache may not work in Windows environments (#22394)ceaf47c4b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
11.8 KiB342 linesraw
1#include "ggml-rpc.h"
2#ifdef _WIN32
3#  define NOMINMAX
4#  define DIRECTORY_SEPARATOR '\\'
5#  include <windows.h>
6#  include <fcntl.h>
7#  include <io.h>
8#else
9#  define DIRECTORY_SEPARATOR '/'
10#  include <unistd.h>
11#  include <sys/stat.h>
12#endif
13#include <algorithm>
14#include <clocale>
15#include <codecvt>
16#include <filesystem>
17#include <regex>
18#include <stdio.h>
19#include <string>
20#include <thread>
21#include <vector>
22
23#if defined(__linux__)
24#include <sys/types.h>
25#include <pwd.h>
26#endif
27
28// NOTE: this is copied from common.cpp to avoid linking with libcommon
29#ifdef _WIN32
30static std::wstring utf8_to_wstring(const std::string & str) {
31    if (str.empty()) {
32        return std::wstring();
33    }
34
35    int size = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), NULL, 0);
36
37    if (size <= 0) {
38        return std::wstring();
39    }
40
41    std::wstring wstr(size, 0);
42    MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), &wstr[0], size);
43
44    return wstr;
45}
46#endif
47
48// NOTE: this is copied from common.cpp to avoid linking with libcommon
49// returns true if successful, false otherwise
50static bool fs_create_directory_with_parents(const std::string & path) {
51#ifdef _WIN32
52    std::wstring wpath = utf8_to_wstring(path);
53
54    // if the path already exists, check whether it's a directory
55    const DWORD attributes = GetFileAttributesW(wpath.c_str());
56    if ((attributes != INVALID_FILE_ATTRIBUTES) && (attributes & FILE_ATTRIBUTE_DIRECTORY)) {
57        return true;
58    }
59
60    size_t pos_slash = 0;
61
62    // process path from front to back, procedurally creating directories
63    while ((pos_slash = path.find('\\', pos_slash)) != std::string::npos) {
64        const std::wstring subpath = wpath.substr(0, pos_slash);
65
66        pos_slash += 1;
67
68        // skip the drive letter, in some systems it can return an access denied error
69        if (subpath.length() == 2 && subpath[1] == ':') {
70            continue;
71        }
72
73        const bool success = CreateDirectoryW(subpath.c_str(), NULL);
74
75        if (!success) {
76            const DWORD error = GetLastError();
77
78            // if the path already exists, ensure that it's a directory
79            if (error == ERROR_ALREADY_EXISTS) {
80                const DWORD attributes = GetFileAttributesW(subpath.c_str());
81                if (attributes == INVALID_FILE_ATTRIBUTES || !(attributes & FILE_ATTRIBUTE_DIRECTORY)) {
82                    return false;
83                }
84            } else {
85                return false;
86            }
87        }
88    }
89
90    return true;
91#else
92    // if the path already exists, check whether it's a directory
93    struct stat info;
94    if (stat(path.c_str(), &info) == 0) {
95        return S_ISDIR(info.st_mode);
96    }
97
98    size_t pos_slash = 1; // skip leading slashes for directory creation
99
100    // process path from front to back, procedurally creating directories
101    while ((pos_slash = path.find('/', pos_slash)) != std::string::npos) {
102        const std::string subpath = path.substr(0, pos_slash);
103        struct stat info;
104
105        // if the path already exists, ensure that it's a directory
106        if (stat(subpath.c_str(), &info) == 0) {
107            if (!S_ISDIR(info.st_mode)) {
108                return false;
109            }
110        } else {
111            // create parent directories
112            const int ret = mkdir(subpath.c_str(), 0755);
113            if (ret != 0) {
114                return false;
115            }
116        }
117
118        pos_slash += 1;
119    }
120
121    return true;
122#endif // _WIN32
123}
124
125// NOTE: this is copied from common.cpp to avoid linking with libcommon
126static std::string fs_get_cache_directory() {
127    std::string cache_directory = "";
128    auto ensure_trailing_slash = [](std::string p) {
129        // Make sure to add trailing slash
130        if (p.back() != DIRECTORY_SEPARATOR) {
131            p += DIRECTORY_SEPARATOR;
132        }
133        return p;
134    };
135    if (getenv("LLAMA_CACHE")) {
136        cache_directory = std::getenv("LLAMA_CACHE");
137    } else {
138#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \
139    defined(__OpenBSD__) || defined(__NetBSD__)
140        if (std::getenv("XDG_CACHE_HOME")) {
141            cache_directory = std::getenv("XDG_CACHE_HOME");
142        } else if (std::getenv("HOME")) {
143            cache_directory = std::getenv("HOME") + std::string("/.cache/");
144        } else {
145#if defined(__linux__)
146            /* no $HOME is defined, fallback to getpwuid */
147            struct passwd *pw = getpwuid(getuid());
148            if ((!pw) || (!pw->pw_dir)) {
149                throw std::runtime_error("Failed to find $HOME directory");
150            }
151
152            cache_directory = std::string(pw->pw_dir) + std::string("/.cache/");
153#else /* defined(__linux__) */
154            throw std::runtime_error("Failed to find $HOME directory");
155#endif /* defined(__linux__) */
156        }
157#elif defined(__APPLE__)
158        cache_directory = std::getenv("HOME") + std::string("/Library/Caches/");
159#elif defined(_WIN32)
160        cache_directory = std::getenv("LOCALAPPDATA");
161#elif defined(__EMSCRIPTEN__)
162        GGML_ABORT("not implemented on this platform");
163#else
164#  error Unknown architecture
165#endif
166        cache_directory = ensure_trailing_slash(cache_directory);
167        cache_directory += "llama.cpp";
168    }
169    return ensure_trailing_slash(cache_directory);
170}
171
172struct rpc_server_params {
173    std::string              host        = "127.0.0.1";
174    int                      port        = 50052;
175    bool                     use_cache   = false;
176    int                      n_threads   = std::max(1U, std::thread::hardware_concurrency()/2);
177    std::vector<std::string> devices;
178};
179
180static void print_usage(int /*argc*/, char ** argv, rpc_server_params params) {
181    fprintf(stderr, "Usage: %s [options]\n\n", argv[0]);
182    fprintf(stderr, "options:\n");
183    fprintf(stderr, "  -h, --help                       show this help message and exit\n");
184    fprintf(stderr, "  -t, --threads N                  number of threads for the CPU device (default: %d)\n", params.n_threads);
185    fprintf(stderr, "  -d, --device <dev1,dev2,...>     comma-separated list of devices\n");
186    fprintf(stderr, "  -H, --host HOST                  host to bind to (default: %s)\n", params.host.c_str());
187    fprintf(stderr, "  -p, --port PORT                  port to bind to (default: %d)\n", params.port);
188    fprintf(stderr, "  -c, --cache                      enable local file cache\n");
189    fprintf(stderr, "\n");
190}
191
192static bool rpc_server_params_parse(int argc, char ** argv, rpc_server_params & params) {
193    std::string arg;
194    for (int i = 1; i < argc; i++) {
195        arg = argv[i];
196        if (arg == "-H" || arg == "--host") {
197            if (++i >= argc) {
198                return false;
199            }
200            params.host = argv[i];
201        } else if (arg == "-t" || arg == "--threads") {
202            if (++i >= argc) {
203                return false;
204            }
205            params.n_threads = std::stoi(argv[i]);
206            if (params.n_threads <= 0) {
207                fprintf(stderr, "error: invalid number of threads: %d\n", params.n_threads);
208                return false;
209            }
210        } else if (arg == "-d" || arg == "--device") {
211            if (++i >= argc) {
212                return false;
213            }
214            const std::regex regex{ R"([,/]+)" };
215            std::string dev_str = argv[i];
216            std::sregex_token_iterator iter(dev_str.begin(), dev_str.end(), regex, -1);
217            std::sregex_token_iterator end;
218            for ( ; iter != end; ++iter) {
219                try {
220                    params.devices.push_back(*iter);
221                } catch (const std::exception & ) {
222                    fprintf(stderr, "error: invalid device: %s\n", iter->str().c_str());
223                    return false;
224                }
225            }
226        } else if (arg == "-p" || arg == "--port") {
227            if (++i >= argc) {
228                return false;
229            }
230            params.port = std::stoi(argv[i]);
231            if (params.port <= 0 || params.port > 65535) {
232                return false;
233            }
234        } else if (arg == "-c" || arg == "--cache") {
235            params.use_cache = true;
236        } else if (arg == "-h" || arg == "--help") {
237            print_usage(argc, argv, params);
238            exit(0);
239        } else {
240            fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
241            print_usage(argc, argv, params);
242            exit(0);
243        }
244    }
245    return true;
246}
247
248static std::vector<ggml_backend_dev_t> get_devices(const rpc_server_params & params) {
249    std::vector<ggml_backend_dev_t> devices;
250    if (!params.devices.empty()) {
251        for (auto device : params.devices) {
252            ggml_backend_dev_t dev = ggml_backend_dev_by_name(device.c_str());
253            if (dev) {
254                devices.push_back(dev);
255            } else {
256                fprintf(stderr, "error: unknown device: %s\n", device.c_str());
257                fprintf(stderr, "available devices:\n");
258                for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
259                    auto * dev = ggml_backend_dev_get(i);
260                    size_t free, total;
261                    ggml_backend_dev_memory(dev, &free, &total);
262                    printf("  %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024);
263                }
264                return {};
265            }
266        }
267    }
268
269    // Try non-CPU devices first
270    if (devices.empty()) {
271        for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
272            ggml_backend_dev_t dev = ggml_backend_dev_get(i);
273            if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
274                devices.push_back(dev);
275            }
276        }
277    }
278
279    // If there are no accelerators, fallback to CPU device
280    if (devices.empty()) {
281        ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
282        if (dev) {
283            devices.push_back(dev);
284        }
285    }
286
287    return devices;
288}
289
290int main(int argc, char * argv[]) {
291    std::setlocale(LC_NUMERIC, "C");
292
293    ggml_backend_load_all();
294
295    rpc_server_params params;
296    if (!rpc_server_params_parse(argc, argv, params)) {
297        fprintf(stderr, "Invalid parameters\n");
298        return 1;
299    }
300
301    if (params.host != "127.0.0.1") {
302        fprintf(stderr, "\n");
303        fprintf(stderr, "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
304        fprintf(stderr, "WARNING: Host ('%s') is != '127.0.0.1'\n", params.host.c_str());
305        fprintf(stderr, "         Never expose the RPC server to an open network!\n");
306        fprintf(stderr, "         This is an experimental feature and is not secure!\n");
307        fprintf(stderr, "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
308        fprintf(stderr, "\n");
309    }
310
311    auto devices = get_devices(params);
312    if (devices.empty()) {
313        fprintf(stderr, "No devices found\n");
314        return 1;
315    }
316    std::string endpoint = params.host + ":" + std::to_string(params.port);
317    const char * cache_dir = nullptr;
318    std::string cache_dir_str;
319    if (params.use_cache) {
320        cache_dir_str = fs_get_cache_directory() + "rpc" + DIRECTORY_SEPARATOR;
321        if (!fs_create_directory_with_parents(cache_dir_str)) {
322            fprintf(stderr, "Failed to create cache directory: %s\n", cache_dir_str.c_str());
323            return 1;
324        }
325        cache_dir = cache_dir_str.c_str();
326    }
327
328    ggml_backend_reg_t reg = ggml_backend_reg_by_name("RPC");
329    if (!reg) {
330        fprintf(stderr, "Failed to find RPC backend\n");
331        return 1;
332    }
333
334    auto start_server_fn = (decltype(ggml_backend_rpc_start_server)*) ggml_backend_reg_get_proc_address(reg, "ggml_backend_rpc_start_server");
335    if (!start_server_fn) {
336        fprintf(stderr, "Failed to obtain RPC backend start server function\n");
337        return 1;
338    }
339
340    start_server_fn(endpoint.c_str(), cache_dir, params.n_threads, devices.size(), devices.data());
341    return 0;
342}