KBaba7/llama.cpp

1#include "ggml-cpu.h"2 3#ifdef GGML_USE_CUDA4#include "ggml-cuda.h"5#endif6 7#ifdef GGML_USE_METAL8#include "ggml-metal.h"9#endif10 11#ifdef GGML_USE_VULKAN12#include "ggml-vulkan.h"13#endif14 15#ifdef GGML_USE_SYCL16#include "ggml-sycl.h"17#endif18 19#include "ggml-rpc.h"20#ifdef _WIN3221#  include <windows.h>22#else23#  include <unistd.h>24#endif25#include <string>26#include <stdio.h>27 28struct rpc_server_params {29    std::string host        = "127.0.0.1";30    int         port        = 50052;31    size_t      backend_mem = 0;32};33 34static void print_usage(int /*argc*/, char ** argv, rpc_server_params params) {35    fprintf(stderr, "Usage: %s [options]\n\n", argv[0]);36    fprintf(stderr, "options:\n");37    fprintf(stderr, "  -h, --help            show this help message and exit\n");38    fprintf(stderr, "  -H HOST, --host HOST  host to bind to (default: %s)\n", params.host.c_str());39    fprintf(stderr, "  -p PORT, --port PORT  port to bind to (default: %d)\n", params.port);40    fprintf(stderr, "  -m MEM, --mem MEM     backend memory size (in MB)\n");41    fprintf(stderr, "\n");42}43 44static bool rpc_server_params_parse(int argc, char ** argv, rpc_server_params & params) {45    std::string arg;46    for (int i = 1; i < argc; i++) {47        arg = argv[i];48        if (arg == "-H" || arg == "--host") {49            if (++i >= argc) {50                return false;51            }52            params.host = argv[i];53        } else if (arg == "-p" || arg == "--port") {54            if (++i >= argc) {55                return false;56            }57            params.port = std::stoi(argv[i]);58            if (params.port <= 0 || params.port > 65535) {59                return false;60            }61        } else if (arg == "-m" || arg == "--mem") {62            if (++i >= argc) {63                return false;64            }65            params.backend_mem = std::stoul(argv[i]) * 1024 * 1024;66        } else if (arg == "-h" || arg == "--help") {67            print_usage(argc, argv, params);68            exit(0);69        } else {70            fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());71            print_usage(argc, argv, params);72            exit(0);73        }74    }75    return true;76}77 78static ggml_backend_t create_backend() {79    ggml_backend_t backend = NULL;80#ifdef GGML_USE_CUDA81    fprintf(stderr, "%s: using CUDA backend\n", __func__);82    backend = ggml_backend_cuda_init(0); // init device 083    if (!backend) {84        fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);85    }86#elif GGML_USE_METAL87    fprintf(stderr, "%s: using Metal backend\n", __func__);88    backend = ggml_backend_metal_init();89    if (!backend) {90        fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);91    }92#elif GGML_USE_VULKAN93    fprintf(stderr, "%s: using Vulkan backend\n", __func__);94    backend = ggml_backend_vk_init(0); // init device 095    if (!backend) {96        fprintf(stderr, "%s: ggml_backend_vulkan_init() failed\n", __func__);97    }98#elif GGML_USE_SYCL99    fprintf(stderr, "%s: using SYCL backend\n", __func__);100    backend = ggml_backend_sycl_init(0); // init device 0101    if (!backend) {102        fprintf(stderr, "%s: ggml_backend_sycl_init() failed\n", __func__);103    }104#endif105 106    // if there aren't GPU Backends fallback to CPU backend107    if (!backend) {108        fprintf(stderr, "%s: using CPU backend\n", __func__);109        backend = ggml_backend_cpu_init();110    }111    return backend;112}113 114static void get_backend_memory(size_t * free_mem, size_t * total_mem) {115#ifdef GGML_USE_CUDA116    ggml_backend_cuda_get_device_memory(0, free_mem, total_mem);117#elif GGML_USE_VULKAN118    ggml_backend_vk_get_device_memory(0, free_mem, total_mem);119#elif GGML_USE_SYCL120    ggml_backend_sycl_get_device_memory(0, free_mem, total_mem);121#else122    #ifdef _WIN32123        MEMORYSTATUSEX status;124        status.dwLength = sizeof(status);125        GlobalMemoryStatusEx(&status);126        *total_mem = status.ullTotalPhys;127        *free_mem = status.ullAvailPhys;128    #else129        long pages = sysconf(_SC_PHYS_PAGES);130        long page_size = sysconf(_SC_PAGE_SIZE);131        *total_mem = pages * page_size;132        *free_mem = *total_mem;133    #endif134#endif135}136 137int main(int argc, char * argv[]) {138    rpc_server_params params;139    if (!rpc_server_params_parse(argc, argv, params)) {140        fprintf(stderr, "Invalid parameters\n");141        return 1;142    }143 144    if (params.host != "127.0.0.1") {145        fprintf(stderr, "\n");146        fprintf(stderr, "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");147        fprintf(stderr, "WARNING: Host ('%s') is != '127.0.0.1'\n", params.host.c_str());148        fprintf(stderr, "         Never expose the RPC server to an open network!\n");149        fprintf(stderr, "         This is an experimental feature and is not secure!\n");150        fprintf(stderr, "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");151        fprintf(stderr, "\n");152    }153 154    ggml_backend_t backend = create_backend();155    if (!backend) {156        fprintf(stderr, "Failed to create backend\n");157        return 1;158    }159    std::string endpoint = params.host + ":" + std::to_string(params.port);160    size_t free_mem, total_mem;161    if (params.backend_mem > 0) {162        free_mem = params.backend_mem;163        total_mem = params.backend_mem;164    } else {165        get_backend_memory(&free_mem, &total_mem);166    }167    printf("Starting RPC server on %s, backend memory: %zu MB\n", endpoint.c_str(), free_mem / (1024 * 1024));168    ggml_backend_rpc_start_server(backend, endpoint.c_str(), free_mem, total_mem);169    ggml_backend_free(backend);170    return 0;171}172