Brunobkr/llama.cpp_AlgMor24_github

1#include "ggml-rpc.h"2#ifdef _WIN323#  define NOMINMAX4#  define DIRECTORY_SEPARATOR '\\'5#  include <windows.h>6#  include <fcntl.h>7#  include <io.h>8#else9#  define DIRECTORY_SEPARATOR '/'10#  include <unistd.h>11#  include <sys/stat.h>12#endif13#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#endif27 28// NOTE: this is copied from common.cpp to avoid linking with libcommon29#ifdef _WIN3230static 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#endif47 48// NOTE: this is copied from common.cpp to avoid linking with libcommon49// returns true if successful, false otherwise50static bool fs_create_directory_with_parents(const std::string & path) {51#ifdef _WIN3252    std::wstring wpath = utf8_to_wstring(path);53 54    // if the path already exists, check whether it's a directory55    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 directories63    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 error69        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 directory79            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#else92    // if the path already exists, check whether it's a directory93    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 creation99 100    // process path from front to back, procedurally creating directories101    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 directory106        if (stat(subpath.c_str(), &info) == 0) {107            if (!S_ISDIR(info.st_mode)) {108                return false;109            }110        } else {111            // create parent directories112            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 // _WIN32123}124 125// NOTE: this is copied from common.cpp to avoid linking with libcommon126static 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 slash130        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#else164#  error Unknown architecture165#endif166        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 first270    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 device280    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}343